Fix ModelListGP.condition_on_observations/fantasize bug (#1250)

Summary:
Pull Request resolved: #1250

`ModelListGP.fantasize` would fail due to `X` being transformed into a list of `X` in `fantasize`, then again being made into a list of `X`s in `condition_on_observations`. This removes the duplication of `X` in condition on observations, using the correctly transformed `X`s for fantasizing.

Fixes #1247

Reviewed By: Balandat

Differential Revision: D36949105

fbshipit-source-id: 3a8a3deb4b78416b8a67cd70863298b009f04209

Author: saitcakmak

botorch/models/model_list_gp_regression.py

@@ -49,15 +49,15 @@ def __init__(self, *gp_models: GPyTorchModel) -> None:
         super().__init__(*gp_models)
 
     def condition_on_observations(
-        self, X: Tensor, Y: Tensor, **kwargs: Any
+        self, X: List[Tensor], Y: Tensor, **kwargs: Any
     ) -> ModelListGP:
         r"""Condition the model on new observations.
 
         Args:
-            X: A `batch_shape x n' x d`-dim Tensor, where `d` is the dimension of
-                the feature space, `n'` is the number of points per batch, and
-                `batch_shape` is the batch shape (must be compatible with the
-                batch shape of the model).
+            X: A `m`-list of `batch_shape x n' x d`-dim Tensors, where `d` is the
+                dimension of the feature space, `n'` is the number of points
+                per batch, and `batch_shape` is the batch shape (must be compatible
+                with the batch shape of the model).
             Y: A `batch_shape' x n' x m`-dim Tensor, where `m` is the number of
                 model outputs, `n'` is the number of points per batch, and
                 `batch_shape'` is the batch shape of the observations.
@@ -73,24 +73,26 @@ def condition_on_observations(
             `n_i + n'` training examples, where the `n'` training examples have
             been added and all test-time caches have been updated.
         """
-        self._validate_tensor_args(
-            X=X, Y=Y, Yvar=kwargs.get("noise", None), strict=False
-        )
-        inputs = [X] * self.num_outputs
         if Y.shape[-1] != self.num_outputs:
             raise BotorchTensorDimensionError(
                 "Incorrect number of outputs for observations. Received "
                 f"{Y.shape[-1]} observation outputs, but model has "
                 f"{self.num_outputs} outputs."
             )
         targets = [Y[..., i] for i in range(Y.shape[-1])]
+        # This should never trigger, posterior call would fail.
+        assert len(targets) == len(X)
         if "noise" in kwargs:
             noise = kwargs.pop("noise")
-            # Note: dimension checks were performed in _validate_tensor_args
+            if noise.shape != Y.shape[-noise.dim() :]:
+                raise BotorchTensorDimensionError(
+                    "The shape of observation noise does not agree with the outcomes. "
+                    f"Received {noise.shape} noise with {Y.shape} outcomes."
+                )
             kwargs_ = {**kwargs, "noise": [noise[..., i] for i in range(Y.shape[-1])]}
         else:
             kwargs_ = kwargs
-        return super().get_fantasy_model(inputs, targets, **kwargs_)
+        return super().get_fantasy_model(X, targets, **kwargs_)
 
     def subset_output(self, idcs: List[int]) -> ModelListGP:
         r"""Subset the model along the output dimension.

test/models/test_gpytorch.py

@@ -215,6 +215,13 @@ def test_validate_tensor_args(self):
                     BotorchTensorDimensionWarning
                 ):
                     GPyTorchModel._validate_tensor_args(X, Y[0], strict=False)
+            # with Yvar
+            if len(output_dim_shape) > 0:
+                Yvar = torch.empty(torch.Size([n]) + output_dim_shape, **tkwargs)
+                GPyTorchModel._validate_tensor_args(X, Y, Yvar)
+                Yvar = torch.empty(n, 5, **tkwargs)
+                with self.assertRaises(BotorchTensorDimensionError):
+                    GPyTorchModel._validate_tensor_args(X, Y, Yvar)
 
     def test_fantasize_flag(self):
         train_X = torch.rand(5, 1)

test/models/test_model_list_gp_regression.py

@@ -17,6 +17,7 @@
 from botorch.models.transforms import Standardize
 from botorch.models.transforms.input import Normalize
 from botorch.posteriors import GPyTorchPosterior
+from botorch.sampling.samplers import IIDNormalSampler
 from botorch.utils.testing import _get_random_data, BotorchTestCase
 from gpytorch.distributions import MultitaskMultivariateNormal, MultivariateNormal
 from gpytorch.kernels import MaternKernel, ScaleKernel
@@ -142,19 +143,19 @@ def test_ModelListGP(self):
             self.assertIsInstance(posterior.mvn, MultivariateNormal)
 
             # test condition_on_observations
-            f_x = torch.rand(2, 1, **tkwargs)
+            f_x = [torch.rand(2, 1, **tkwargs) for _ in range(2)]
             f_y = torch.rand(2, 2, **tkwargs)
             cm = model.condition_on_observations(f_x, f_y)
             self.assertIsInstance(cm, ModelListGP)
 
             # test condition_on_observations batched
-            f_x = torch.rand(3, 2, 1, **tkwargs)
+            f_x = [torch.rand(3, 2, 1, **tkwargs) for _ in range(2)]
             f_y = torch.rand(3, 2, 2, **tkwargs)
             cm = model.condition_on_observations(f_x, f_y)
             self.assertIsInstance(cm, ModelListGP)
 
             # test condition_on_observations batched (fast fantasies)
-            f_x = torch.rand(2, 1, **tkwargs)
+            f_x = [torch.rand(2, 1, **tkwargs) for _ in range(2)]
             f_y = torch.rand(3, 2, 2, **tkwargs)
             cm = model.condition_on_observations(f_x, f_y)
             self.assertIsInstance(cm, ModelListGP)
@@ -163,6 +164,10 @@ def test_ModelListGP(self):
             with self.assertRaises(BotorchTensorDimensionError):
                 model.condition_on_observations(f_x, torch.rand(3, 2, 3, **tkwargs))
 
+            # test X having wrong size
+            with self.assertRaises(AssertionError):
+                cm = model.condition_on_observations(f_x[:1], f_y)
+
             # test posterior transform
             X = torch.rand(3, 1, **tkwargs)
             weights = torch.tensor([1, 2], **tkwargs)
@@ -222,21 +227,21 @@ def test_ModelListGP_fixed_noise(self):
             self.assertIsInstance(posterior.mvn, MultivariateNormal)
 
             # test condition_on_observations
-            f_x = torch.rand(2, 1, **tkwargs)
+            f_x = [torch.rand(2, 1, **tkwargs) for _ in range(2)]
             f_y = torch.rand(2, 2, **tkwargs)
             noise = 0.1 + 0.1 * torch.rand_like(f_y)
             cm = model.condition_on_observations(f_x, f_y, noise=noise)
             self.assertIsInstance(cm, ModelListGP)
 
             # test condition_on_observations batched
-            f_x = torch.rand(3, 2, 1, **tkwargs)
+            f_x = [torch.rand(3, 2, 1, **tkwargs) for _ in range(2)]
             f_y = torch.rand(3, 2, 2, **tkwargs)
             noise = 0.1 + 0.1 * torch.rand_like(f_y)
             cm = model.condition_on_observations(f_x, f_y, noise=noise)
             self.assertIsInstance(cm, ModelListGP)
 
             # test condition_on_observations batched (fast fantasies)
-            f_x = torch.rand(2, 1, **tkwargs)
+            f_x = [torch.rand(2, 1, **tkwargs) for _ in range(2)]
             f_y = torch.rand(3, 2, 2, **tkwargs)
             noise = 0.1 + 0.1 * torch.rand(2, 2, **tkwargs)
             cm = model.condition_on_observations(f_x, f_y, noise=noise)
@@ -295,3 +300,15 @@ def test_transform_revert_train_inputs(self):
             )
             self.assertTrue(m._has_transformed_inputs)
             self.assertTrue(torch.equal(m._original_train_inputs, org_inputs[i]))
+
+    def test_fantasize(self):
+        m1 = SingleTaskGP(torch.rand(5, 2), torch.rand(5, 1)).eval()
+        m2 = SingleTaskGP(torch.rand(5, 2), torch.rand(5, 1)).eval()
+        modellist = ModelListGP(m1, m2)
+        fm = modellist.fantasize(torch.rand(3, 2), sampler=IIDNormalSampler(2))
+        self.assertIsInstance(fm, ModelListGP)
+        for i in range(2):
+            fm_i = fm.models[i]
+            self.assertIsInstance(fm_i, SingleTaskGP)
+            self.assertEqual(fm_i.train_inputs[0].shape, torch.Size([2, 8, 2]))
+            self.assertEqual(fm_i.train_targets.shape, torch.Size([2, 8]))