Skip tensor validation in condition_on_observations when fantasizing (#2378)

Summary:
Pull Request resolved: #2378

Resolves #2376

When fantasizing, we expect `Y` to have an additional batch dimension, so it will fail the tensor shape checks. Since `Y` is produced by the model itself, it should always be a valid input. Disabling these checks to avoid misleading warnings.

Reviewed By: Balandat

Differential Revision: D58564318

fbshipit-source-id: 5d21784b116a7a72c93d3904dcdd36e172e785ad

Author: saitcakmak

botorch/models/gpytorch.py

@@ -39,6 +39,7 @@
     mod_batch_shape,
     multioutput_to_batch_mode_transform,
 )
+from botorch.models.utils.assorted import fantasize as fantasize_flag
 from botorch.posteriors.fully_bayesian import GaussianMixturePosterior
 from botorch.posteriors.gpytorch import GPyTorchPosterior
 from botorch.utils.multitask import separate_mtmvn
@@ -246,9 +247,11 @@ def condition_on_observations(
             if not isinstance(self, BatchedMultiOutputGPyTorchModel):
                 # `noise` is assumed to already be outcome-transformed.
                 Y, _ = self.outcome_transform(Y=Y, Yvar=Yvar)
-        # validate using strict=False, since we cannot tell if Y has an explicit
-        # output dimension
-        self._validate_tensor_args(X=X, Y=Y, Yvar=Yvar, strict=False)
+        # Validate using strict=False, since we cannot tell if Y has an explicit
+        # output dimension. Do not check shapes when fantasizing as they are
+        # not expected to match.
+        if fantasize_flag.off():
+            self._validate_tensor_args(X=X, Y=Y, Yvar=Yvar, strict=False)
         if Y.size(-1) == 1:
             Y = Y.squeeze(-1)
             if Yvar is not None:
@@ -505,7 +508,9 @@ def condition_on_observations(
             # We need to apply transforms before shifting batch indices around.
             # `noise` is assumed to already be outcome-transformed.
             Y, _ = self.outcome_transform(Y)
-        self._validate_tensor_args(X=X, Y=Y, Yvar=noise, strict=False)
+        # Do not check shapes when fantasizing as they are not expected to match.
+        if fantasize_flag.off():
+            self._validate_tensor_args(X=X, Y=Y, Yvar=noise, strict=False)
         inputs = X
         if self._num_outputs > 1:
             inputs, targets, noise = multioutput_to_batch_mode_transform(

botorch/models/higher_order_gp.py

@@ -25,6 +25,7 @@
 from botorch.models.transforms.input import InputTransform
 from botorch.models.transforms.outcome import OutcomeTransform, Standardize
 from botorch.models.utils import gpt_posterior_settings
+from botorch.models.utils.assorted import fantasize as fantasize_flag
 from botorch.models.utils.gpytorch_modules import (
     get_gaussian_likelihood_with_gamma_prior,
 )
@@ -414,7 +415,9 @@ def condition_on_observations(
         if hasattr(self, "outcome_transform"):
             # we need to apply transforms before shifting batch indices around
             Y, noise = self.outcome_transform(Y=Y, Yvar=noise)
-        self._validate_tensor_args(X=X, Y=Y, Yvar=noise, strict=False)
+        # Do not check shapes when fantasizing as they are not expected to match.
+        if fantasize_flag.off():
+            self._validate_tensor_args(X=X, Y=Y, Yvar=noise, strict=False)
 
         # we don't need to do un-squeezing because Y already is batched
         # we don't support fixed noise here yet

test/models/test_gpytorch.py

@@ -250,6 +250,20 @@ def test_validate_tensor_args(self) -> None:
                     ):
                         GPyTorchModel._validate_tensor_args(X, Y, Yvar, strict=strict)
 
+    def test_condition_on_observations_tensor_validation(self) -> None:
+        model = SimpleGPyTorchModel(torch.rand(5, 1), torch.randn(5, 1))
+        model.posterior(torch.rand(2, 1))  # evaluate the model to form caches.
+        # Outside of fantasize, the inputs are validated.
+        with self.assertWarnsRegex(
+            BotorchTensorDimensionWarning, "Non-strict enforcement of"
+        ):
+            model.condition_on_observations(torch.randn(2, 1), torch.randn(5, 2, 1))
+        # Inside of fantasize, the inputs are not validated.
+        with fantasize(), warnings.catch_warnings(record=True) as ws:
+            warnings.filterwarnings("always", category=BotorchTensorDimensionWarning)
+            model.condition_on_observations(torch.randn(2, 1), torch.randn(5, 2, 1))
+        self.assertFalse(any(w.category is BotorchTensorDimensionWarning for w in ws))
+
     def test_fantasize_flag(self):
         train_X = torch.rand(5, 1)
         train_Y = torch.sin(train_X)