Fix RFF bug when using FixedNoiseGP models (#1528)

Summary:
Pull Request resolved: #1528

`_model.likelihood.noise` is a `... x d`-dim tensor for fixed noise models with `d`-dimensional `train_X`. This leads to errors when constructing the RFF. This updates the `sigma_sq` to use the mean across the input dimensions, which was used prior to #1336.

Reviewed By: Balandat

Differential Revision: D41568680

fbshipit-source-id: fef740ebd62e49d63397ebc4ea1d1dc883f798aa

Author: saitcakmak

botorch/utils/gp_sampling.py

@@ -462,7 +462,7 @@ def get_gp_samples(
         phi_X = basis(train_X)
         # Sample weights from bayesian linear model.
         # weights.sample().shape == (n_samples, batch_shape_input, num_rff_features)
-        sigma_sq = _model.likelihood.noise
+        sigma_sq = _model.likelihood.noise.mean(dim=-1, keepdim=True)
         if len(basis.kernel_batch_shape) > 0:
             sigma_sq = sigma_sq.unsqueeze(-2)
         mvn = get_weights_posterior(

test/utils/test_gp_sampling.py

@@ -11,7 +11,7 @@
 import torch
 from botorch.models.converter import batched_to_model_list
 from botorch.models.deterministic import DeterministicModel
-from botorch.models.gp_regression import SingleTaskGP
+from botorch.models.gp_regression import FixedNoiseGP, SingleTaskGP
 from botorch.models.model import ModelList
 from botorch.models.multitask import MultiTaskGP
 from botorch.models.transforms.input import Normalize
@@ -644,3 +644,20 @@ def test_get_gp_samples(self):
                     expected = torch.Size([13, 5, 3, m])
                 Y_batched = gp_samples.posterior(test_X).mean
                 self.assertEqual(Y_batched.shape, expected)
+
+    def test_with_fixed_noise(self):
+        for n_samples in (1, 20):
+            gp_samples = get_gp_samples(
+                model=FixedNoiseGP(
+                    torch.rand(5, 3, dtype=torch.double),
+                    torch.randn(5, 1, dtype=torch.double),
+                    torch.rand(5, 1, dtype=torch.double) * 0.1,
+                ),
+                num_outputs=1,
+                n_samples=n_samples,
+            )
+            samples = gp_samples(torch.rand(2, 3))
+            expected_shape = (
+                torch.Size([2, 1]) if n_samples == 1 else torch.Size([n_samples, 2, 1])
+            )
+            self.assertEqual(samples.shape, expected_shape)