fix fantasization with known noise and empty X (#2090)

Summary: Pull Request resolved: #2090 Previously with a FixedNoiseGaussianLikelihood `noise` was not passed as an argument when X is empty (as with decoupled AFs). Reviewed By: Balandat Differential Revision: D51115939 fbshipit-source-id: c979afeb491d9428ea0c03a1151b2a7ab78c6c44
pytorch · Nov 8, 2023 · 4677556 · 4677556
1 parent 4b30772
commit 4677556
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 23 deletions.
diff --git a/botorch/models/model.py b/botorch/models/model.py
@@ -374,6 +374,8 @@ def fantasize(
                 + torch.Size([0, self.num_outputs])
             )
             Y = torch.empty(output_shape, dtype=X.dtype, device=X.device)
+            if observation_noise is not None:
+                kwargs["noise"] = observation_noise.expand(Y.shape[1:])
             return self.condition_on_observations(
                 X=self.transform_inputs(X),
                 Y=Y,

diff --git a/test/models/test_model_list_gp_regression.py b/test/models/test_model_list_gp_regression.py
@@ -430,31 +430,48 @@ def test_fantasize(self):
                     (3, 2), 0.3, dtype=x1.dtype, device=x1.device
                 )
                 observation_noise[:, 1] = 0.4
-                fm = modellist.fantasize(
-                    torch.rand(3, 2),
-                    sampler=ListSampler(sampler1, sampler2),
-                    evaluation_mask=eval_mask,
-                    observation_noise=observation_noise,
-                )
-                self.assertIsInstance(fm, ModelListGP)
-                for i in range(2):
-                    fm_i = fm.models[i]
-                    self.assertIsInstance(fm_i, SingleTaskGP)
-                    self.assertIsInstance(fm_i.likelihood, FixedNoiseGaussianLikelihood)
-                    num_points = 7 - i
-                    self.assertEqual(
-                        fm_i.train_inputs[0].shape, torch.Size([2, num_points, 2])
-                    )
-                    self.assertEqual(
-                        fm_i.train_targets.shape, torch.Size([2, num_points])
+                for obs_noise in (None, observation_noise):
+                    fm = modellist.fantasize(
+                        torch.rand(3, 2),
+                        sampler=ListSampler(sampler1, sampler2),
+                        evaluation_mask=eval_mask,
+                        observation_noise=obs_noise,
                     )
-                    # check observation_noise
-                    self.assertTrue(
-                        torch.equal(
-                            fm_i.likelihood.noise[..., -num_designs_per_output[i] :],
-                            observation_noise[-num_designs_per_output[i] :, i],
+                    self.assertIsInstance(fm, ModelListGP)
+                    for i in range(2):
+                        fm_i = fm.models[i]
+                        self.assertIsInstance(fm_i, SingleTaskGP)
+                        self.assertIsInstance(
+                            fm_i.likelihood, FixedNoiseGaussianLikelihood
                         )
-                    )
+                        num_points = 7 - i
+                        self.assertEqual(
+                            fm_i.train_inputs[0].shape, torch.Size([2, num_points, 2])
+                        )
+                        self.assertEqual(
+                            fm_i.train_targets.shape, torch.Size([2, num_points])
+                        )
+                        # check observation_noise
+                        if obs_noise is not None:
+                            self.assertTrue(
+                                torch.equal(
+                                    fm_i.likelihood.noise[
+                                        ..., -num_designs_per_output[i] :
+                                    ],
+                                    observation_noise[-num_designs_per_output[i] :, i],
+                                )
+                            )
+                        else:
+                            self.assertTrue(
+                                torch.allclose(
+                                    fm_i.likelihood.noise[
+                                        ..., -num_designs_per_output[i] :
+                                    ],
+                                    modellist.models[i]
+                                    .likelihood.noise[..., -num_designs_per_output[i] :]
+                                    .mean(),
+                                )
+                            )
 
     def test_fantasize_with_outcome_transform(self) -> None:
         """