Properly pass noise levels to GPs with FixedNoiseGaussianLikelihood

botorch/models/gpytorch.py

@@ -17,6 +17,7 @@
 from gpytorch import settings as gpt_settings
 from gpytorch.distributions import MultitaskMultivariateNormal, MultivariateNormal
 from gpytorch.lazy import lazify
+from gpytorch.likelihoods.gaussian_likelihood import FixedNoiseGaussianLikelihood
 from torch import Tensor
 
 from .. import settings
@@ -181,7 +182,12 @@ def posterior(
                 )
             mvn = self(X)
             if observation_noise:
-                mvn = self.likelihood(mvn, X)
+                if isinstance(self.likelihood, FixedNoiseGaussianLikelihood):
+                    # Use the mean of the previous noise values (TODO: be smarter here).
+                    noise = self.likelihood.noise.mean().expand(X.shape[:-1])
+                    mvn = self.likelihood(mvn, X, noise=noise)
+                else:
+                    mvn = self.likelihood(mvn, X)
             if self._num_outputs > 1:
                 mean_x = mvn.mean
                 covar_x = mvn.covariance_matrix
@@ -295,9 +301,15 @@ def posterior(
             if output_indices is not None:
                 mvns = [self.forward_i(i, X) for i in output_indices]
                 if observation_noise:
+                    lh_kwargs = [
+                        {"noise": lh.noise.mean().expand(X.shape[:-1])}
+                        if isinstance(lh, FixedNoiseGaussianLikelihood)
+                        else {}
+                        for lh in self.likelihood.likelihoods
+                    ]
                     mvns = [
-                        self.likelihood_i(i, mvn, X)
-                        for i, mvn in zip(output_indices, mvns)
+                        self.likelihood_i(i, mvn, X, **lkws)
+                        for i, mvn, lkws in zip(output_indices, mvns, lh_kwargs)
                     ]
             else:
                 mvns = self(*[X for _ in range(self.num_outputs)])

test/models/fidelity/test_gp_regression_fidelity.py

@@ -4,10 +4,12 @@
 
 import math
 import unittest
+import warnings
 
 import torch
 from botorch import fit_gpytorch_model
-from botorch.exceptions import UnsupportedError
+from botorch.exceptions.errors import UnsupportedError
+from botorch.exceptions.warnings import OptimizationWarning
 from botorch.models.fidelity.gp_regression_fidelity import (
     SingleTaskGPLTKernel,
     SingleTaskMultiFidelityGP,
@@ -118,9 +120,13 @@ def test_gp(self, cuda=False):
                         mll = ExactMarginalLogLikelihood(model.likelihood, model).to(
                             **tkwargs
                         )
-                        fit_gpytorch_model(
-                            mll, sequential=False, options={"maxiter": 1}
-                        )
+                        with warnings.catch_warnings():
+                            warnings.filterwarnings(
+                                "ignore", category=OptimizationWarning
+                            )
+                            fit_gpytorch_model(
+                                mll, sequential=False, options={"maxiter": 1}
+                            )
 
                         # test init
                         self.assertIsInstance(model.mean_module, ConstantMean)

test/models/test_gp_regression.py

@@ -351,10 +351,13 @@ def test_fantasize(self, cuda=False):
 
 
 def _get_pvar_expected(posterior, model, X, num_outputs):
+    lh_kwargs = {}
+    if isinstance(model.likelihood, FixedNoiseGaussianLikelihood):
+        lh_kwargs["noise"] = model.likelihood.noise.mean().expand(X.shape[:-1])
     if num_outputs == 1:
-        return model.likelihood(posterior.mvn, X).variance.unsqueeze(-1)
+        return model.likelihood(posterior.mvn, X, **lh_kwargs).variance.unsqueeze(-1)
     X_, odi = add_output_dim(X=X, original_batch_shape=model._input_batch_shape)
-    pvar_exp = model.likelihood(model(X_), X_).variance
+    pvar_exp = model.likelihood(model(X_), X_, **lh_kwargs).variance
     return torch.stack(
         [pvar_exp.select(dim=odi, index=i) for i in range(num_outputs)], dim=-1
     )

test/test_fit.py

@@ -8,7 +8,7 @@
 
 import torch
 from botorch import fit_gpytorch_model
-from botorch.exceptions.warnings import OptimizationWarning
+from botorch.exceptions.warnings import BotorchWarning, OptimizationWarning
 from botorch.models import FixedNoiseGP, HeteroskedasticSingleTaskGP, SingleTaskGP
 from botorch.optim.fit import (
     OptimizationIteration,
@@ -151,7 +151,11 @@ def test_fit_gpytorch_model_singular(self, cuda=False):
             mll = ExactMarginalLogLikelihood(gp.likelihood, gp)
             mll.to(device=device, dtype=dtype)
             # this will do multiple retries (and emit warnings, which is desired)
-            fit_gpytorch_model(mll, options=options, max_retries=2)
+            with warnings.catch_warnings(record=True) as ws:
+                fit_gpytorch_model(mll, options=options, max_retries=2)
+                self.assertTrue(
+                    any(issubclass(w.category, OptimizationWarning) for w in ws)
+                )
 
     def test_fit_gpytorch_model_singular_cuda(self):
         if torch.cuda.is_available():
@@ -168,8 +172,7 @@ def test_fit_gpytorch_model_sequential(self, cuda=False):
         options = {"disp": False, "maxiter": 1}
         for double in (False, True):
             for kind in ("SingleTaskGP", "FixedNoiseGP", "HeteroskedasticSingleTaskGP"):
-                with warnings.catch_warnings():
-                    warnings.filterwarnings("ignore", category=OptimizationWarning)
+                with warnings.catch_warnings(record=True) as ws:
                     mll = self._getBatchedModel(kind=kind, double=double, cuda=cuda)
                     mll = fit_gpytorch_model(mll, options=options, max_retries=1)
                     mll = self._getBatchedModel(kind=kind, double=double, cuda=cuda)
@@ -180,6 +183,17 @@ def test_fit_gpytorch_model_sequential(self, cuda=False):
                     mll = fit_gpytorch_model(
                         mll, options=options, sequential=False, max_retries=1
                     )
+                    if kind == "HeteroskedasticSingleTaskGP":
+                        self.assertTrue(
+                            any(issubclass(w.category, BotorchWarning) for w in ws)
+                        )
+                        self.assertTrue(
+                            any(
+                                "Failed to convert ModelList to batched model"
+                                in str(w.message)
+                                for w in ws
+                            )
+                        )
 
     def test_fit_gpytorch_model_sequential_cuda(self):
         if torch.cuda.is_available():