Repository: cornellius-gp/gpytorch
Pull request #636 (2019-04-10 21:28:41)

Add support for fixed noise fantasies
This allows to use `get_fantasy_model` with `FixedNoiseGaussianLikelihood` if specifying additional noise levels for the fantasy points using a kwarg to `get_fantasy_model`.

Will do more thorough testing and add a unit test.


Patch:
diff --git a/gpytorch/likelihoods/gaussian_likelihood.py b/gpytorch/likelihoods/gaussian_likelihood.py
index 8cecfb6d..ffae0e89 100644
--- a/gpytorch/likelihoods/gaussian_likelihood.py
+++ b/gpytorch/likelihoods/gaussian_likelihood.py
@@ -1,5 +1,7 @@
 #!/usr/bin/env python3
 
+from copy import deepcopy
+
 import math
 import warnings
 from typing import Any, Optional
@@ -37,8 +39,8 @@ class _GaussianLikelihoodBase(Likelihood):
         return self.noise_covar(*params, shape=shape, **kwargs)
 
     def forward(self, function_samples: Tensor, *params: Any, **kwargs: Any) -> base_distributions.Normal:
-        observation_noise = self._shaped_noise_covar(function_samples.shape, *params, **kwargs).diag()
-        return base_distributions.Normal(function_samples, observation_noise.sqrt())
+        noise = self._shaped_noise_covar(function_samples.shape, *params, **kwargs).diag()
+        return base_distributions.Normal(function_samples, noise.sqrt())
 
     def marginal(self, function_dist: MultivariateNormal, *params: Any, **kwargs: Any) -> MultivariateNormal:
         mean, covar = function_dist.mean, function_dist.lazy_covariance_matrix
@@ -94,9 +96,8 @@ class FixedNoiseGaussianLikelihood(_GaussianLikelihoodBase):
         :attr:`learn_additional_noise` (bool, optional):
             Set to true if you additionally want to learn added diagonal noise, similar to GaussianLikelihood.
 
-    Note that this likelihood takes an additional argument when you call it, `observation_noise`, that adds
-    a specified amount of noise to the passed MultivariateNormal. This allows for adding known observational noise
-    to test data.
+    Note that this likelihood takes an additional argument when you call it, `noise`, that adds a specified amount
+    of noise to the passed MultivariateNormal. This allows for adding known observational noise to test data.
 
     Example:
         >>> train_x = torch.randn(55, 2)
@@ -106,7 +107,7 @@ class FixedNoiseGaussianLikelihood(_GaussianLikelihoodBase):
         >>>
         >>> test_x = torch.randn(21, 2)
         >>> test_noises = torch.ones(21) * 0.02
-        >>> pred_y = likelihood(gp_model(test_x), observation_noise=test_noises)
+        >>> pred_y = likelihood(gp_model(test_x), noise=test_noises)
     """
     def __init__(
         self,
@@ -155,6 +156,21 @@ class FixedNoiseGaussianLikelihood(_GaussianLikelihoodBase):
             )
         self.second_noise_covar.initialize(noise=value)
 
+    def get_fantasy_likelihood(self, **kwargs):
+        if "noise" not in kwargs:
+            raise RuntimeError("FixedNoiseGaussianLikelihood.fantasize requires a `noise` kwarg")
+        old_noise_covar = self.noise_covar
+        self.noise_covar = None
+        fantasy_liklihood = deepcopy(self)
+        self.noise_covar = old_noise_covar
+
+        old_noise = old_noise_covar.noise
+        new_noise = kwargs.get("noise")
+        if old_noise.dim() != new_noise.dim():
+            old_noise = old_noise.expand(*new_noise.shape[:-1], old_noise.shape[-1])
+        fantasy_liklihood.noise_covar = FixedGaussianNoise(noise=torch.cat([old_noise, new_noise], -1))
+        return fantasy_liklihood
+
     def _shaped_noise_covar(self, base_shape: torch.Size, *params: Any, **kwargs: Any):
         if len(params) > 0:
             # we can infer the shape from the params
@@ -170,7 +186,7 @@ class FixedNoiseGaussianLikelihood(_GaussianLikelihoodBase):
         elif isinstance(res, ZeroLazyTensor):
             warnings.warn(
                 "You have passed data through a FixedNoiseGaussianLikelihood that did not match the size "
-                "of the fixed noise, *and* you did not specify observation_noise. This is treated as a no-op."
+                "of the fixed noise, *and* you did not specify noise. This is treated as a no-op."
             )
 
         return res
diff --git a/gpytorch/likelihoods/likelihood.py b/gpytorch/likelihoods/likelihood.py
index c45289a2..cc51aa13 100644
--- a/gpytorch/likelihoods/likelihood.py
+++ b/gpytorch/likelihoods/likelihood.py
@@ -1,5 +1,7 @@
 #!/usr/bin/env python3
 
+from copy import deepcopy
+
 import functools
 import torch
 import warnings
@@ -109,6 +111,9 @@ class _Likelihood(Module, _ClassWithDeprecatedBatchSize):
         )
         return self.expected_log_prob(observations, function_dist)
 
+    def get_fantasy_likelihood(self, **kwargs):
+        return deepcopy(self)
+
     def __call__(self, input, *params, **kwargs):
         # Conditional
         if torch.is_tensor(input):
diff --git a/gpytorch/likelihoods/multitask_gaussian_likelihood.py b/gpytorch/likelihoods/multitask_gaussian_likelihood.py
index e0160027..7adac673 100644
--- a/gpytorch/likelihoods/multitask_gaussian_likelihood.py
+++ b/gpytorch/likelihoods/multitask_gaussian_likelihood.py
@@ -117,9 +117,9 @@ class _MultitaskGaussianLikelihoodBase(_GaussianLikelihoodBase):
         return task_covar
 
     def forward(self, function_samples: Tensor, *params: Any, **kwargs: Any) -> base_distributions.Normal:
-        observation_noise = self._shaped_noise_covar(function_samples.shape, *params, **kwargs).diag()
-        observation_noise = observation_noise.view(*observation_noise.shape[:-1], *function_samples.shape[-2:])
-        return base_distributions.Normal(function_samples, observation_noise.sqrt())
+        noise = self._shaped_noise_covar(function_samples.shape, *params, **kwargs).diag()
+        noise = noise.view(*noise.shape[:-1], *function_samples.shape[-2:])
+        return base_distributions.Normal(function_samples, noise.sqrt())
 
 
 class MultitaskGaussianLikelihood(_MultitaskGaussianLikelihoodBase):
diff --git a/gpytorch/likelihoods/noise_models.py b/gpytorch/likelihoods/noise_models.py
index 9b9e2eae..22a07141 100644
--- a/gpytorch/likelihoods/noise_models.py
+++ b/gpytorch/likelihoods/noise_models.py
@@ -18,14 +18,27 @@ class Noise(Module):
 
 
 class _HomoskedasticNoiseBase(Noise):
-    def __init__(self, noise_prior=None, noise_constraint=None, batch_shape=torch.Size(), num_tasks=1):
+    def __init__(
+        self,
+        noise_prior=None,
+        noise_constraint=None,
+        batch_shape=torch.Size(),
+        num_tasks=1,
+    ):
         super().__init__()
         if noise_constraint is None:
             noise_constraint = GreaterThan(1e-4)
 
-        self.register_parameter(name="raw_noise", parameter=Parameter(torch.zeros(*batch_shape, num_tasks)))
+        self.register_parameter(
+            name="raw_noise", parameter=Parameter(torch.zeros(*batch_shape, num_tasks))
+        )
         if noise_prior is not None:
-            self.register_prior("noise_prior", noise_prior, lambda: self.noise, lambda v: self._set_noise(v))
+            self.register_prior(
+                "noise_prior",
+                noise_prior,
+                lambda: self.noise,
+                lambda v: self._set_noise(v),
+            )
 
         self.register_constraint("raw_noise", noise_constraint)
 
@@ -42,7 +55,9 @@ class _HomoskedasticNoiseBase(Noise):
             value = torch.as_tensor(value).to(self.raw_noise)
         self.initialize(raw_noise=self.raw_noise_constraint.inverse_transform(value))
 
-    def forward(self, *params: Any, shape: Optional[torch.Size] = None) -> DiagLazyTensor:
+    def forward(
+        self, *params: Any, shape: Optional[torch.Size] = None
+    ) -> DiagLazyTensor:
         """In the homoskedastic case, the parameters are only used to infer the required shape.
         Here are the possible scenarios:
         - non-batched noise, non-batched input, non-MT -> noise_diag shape is `n`
@@ -75,7 +90,9 @@ class _HomoskedasticNoiseBase(Noise):
 
 
 class HomoskedasticNoise(_HomoskedasticNoiseBase):
-    def __init__(self, noise_prior=None, noise_constraint=None, batch_shape=torch.Size()):
+    def __init__(
+        self, noise_prior=None, noise_constraint=None, batch_shape=torch.Size()
+    ):
         super().__init__(
             noise_prior=noise_prior,
             noise_constraint=noise_constraint,
@@ -85,7 +102,13 @@ class HomoskedasticNoise(_HomoskedasticNoiseBase):
 
 
 class MultitaskHomoskedasticNoise(_HomoskedasticNoiseBase):
-    def __init__(self, num_tasks, noise_prior=None, noise_constraint=None, batch_shape=torch.Size()):
+    def __init__(
+        self,
+        num_tasks,
+        noise_prior=None,
+        noise_constraint=None,
+        batch_shape=torch.Size(),
+    ):
         super().__init__(
             noise_prior=noise_prior,
             noise_constraint=noise_constraint,
@@ -104,17 +127,26 @@ class HeteroskedasticNoise(Noise):
         self._noise_indices = noise_indices
 
     def forward(
-        self, *params: Any, batch_shape: Optional[torch.Size] = None, shape: Optional[torch.Size] = None
+        self,
+        *params: Any,
+        batch_shape: Optional[torch.Size] = None,
+        shape: Optional[torch.Size] = None
     ) -> DiagLazyTensor:
         if len(params) == 1 and not torch.is_tensor(params[0]):
             output = self.noise_model(*params[0])
         else:
             output = self.noise_model(*params)
         if not isinstance(output, MultivariateNormal):
-            raise NotImplementedError("Currently only noise models that return a MultivariateNormal are supported")
+            raise NotImplementedError(
+                "Currently only noise models that return a MultivariateNormal are supported"
+            )
         # note: this also works with MultitaskMultivariateNormal, where this
         # will return a batched DiagLazyTensors of size n x num_tasks x num_tasks
-        noise_diag = output.mean if self._noise_indices is None else output.mean[..., self._noise_indices]
+        noise_diag = (
+            output.mean
+            if self._noise_indices is None
+            else output.mean[..., self._noise_indices]
+        )
         return DiagLazyTensor(self._noise_constraint.transform(noise_diag))
 
 
@@ -127,15 +159,15 @@ class FixedGaussianNoise(Module):
         self,
         *params: Any,
         shape: Optional[torch.Size] = None,
-        observation_noise: Optional[Tensor] = None,
+        noise: Optional[Tensor] = None,
         **kwargs: Any
     ) -> DiagLazyTensor:
         if shape is None:
             p = params[0] if torch.is_tensor(params[0]) else params[0][0]
             shape = p.shape if len(p.shape) == 1 else p.shape[:-1]
 
-        if observation_noise is not None:
-            return DiagLazyTensor(observation_noise)
+        if noise is not None:
+            return DiagLazyTensor(noise)
         elif shape[-1] == self.noise.shape[-1]:
             return DiagLazyTensor(self.noise)
         else:
diff --git a/gpytorch/models/exact_gp.py b/gpytorch/models/exact_gp.py
index abbb1864..5f2679ce 100644
--- a/gpytorch/models/exact_gp.py
+++ b/gpytorch/models/exact_gp.py
@@ -114,7 +114,7 @@ class ExactGP(GP):
         if not isinstance(inputs, list):
             inputs = [inputs]
 
-        inputs = list(i.unsqueeze(-1) if i.ndimension() == 1 else i for i in inputs)
+        inputs = [i.unsqueeze(-1) if i.ndimension() == 1 else i for i in inputs]
 
         # If input is n x d but targets is b x n x d, expand input to b x n x d
         for i, input in enumerate(inputs):
@@ -133,28 +133,38 @@ class ExactGP(GP):
         full_inputs = [torch.cat([train_input, input], dim=-2) for train_input, input in zip(train_inputs, inputs)]
         full_targets = torch.cat([train_targets, targets], dim=-1)
 
+        try:
+            fantasy_kwargs = {"noise": kwargs.pop("noise")}
+        except KeyError:
+            fantasy_kwargs = {}
+
         full_output = super(ExactGP, self).__call__(*full_inputs, **kwargs)
 
         # Copy model without copying training data or prediction strategy (since we'll overwrite those)
         old_pred_strat = self.prediction_strategy
         old_train_inputs = self.train_inputs
         old_train_targets = self.train_targets
+        old_likelihood = self.likelihood
         self.prediction_strategy = None
         self.train_inputs = None
         self.train_targets = None
+        self.likelihood = None
         new_model = deepcopy(self)
         self.prediction_strategy = old_pred_strat
         self.train_inputs = old_train_inputs
         self.train_targets = old_train_targets
+        self.likelihood = old_likelihood
 
         new_model.train_inputs = full_inputs
         new_model.train_targets = full_targets
+        new_model.likelihood = old_likelihood.get_fantasy_likelihood(**fantasy_kwargs)
         new_model.prediction_strategy = old_pred_strat.get_fantasy_strategy(
             inputs,
             targets,
             full_inputs,
             full_targets,
-            full_output
+            full_output,
+            **fantasy_kwargs,
         )
 
         return new_model
@@ -179,7 +189,7 @@ class ExactGP(GP):
 
     def __call__(self, *args, **kwargs):
         train_inputs = list(self.train_inputs) if self.train_inputs is not None else []
-        inputs = list(i.unsqueeze(-1) if i.ndimension() == 1 else i for i in args)
+        inputs = [i.unsqueeze(-1) if i.ndimension() == 1 else i for i in args]
 
         # Training mode: optimizing
         if self.training:
diff --git a/gpytorch/models/exact_prediction_strategies.py b/gpytorch/models/exact_prediction_strategies.py
index 22ee1704..08034146 100644
--- a/gpytorch/models/exact_prediction_strategies.py
+++ b/gpytorch/models/exact_prediction_strategies.py
@@ -41,7 +41,7 @@ class DefaultPredictionStrategy(object):
         self.lik_train_train_covar = mvn.lazy_covariance_matrix
 
     def __deepcopy__(self, memo):
-        # deepcopying prediciton strategies of a model evaluated on inputs that require gradients fails
+        # deepcopying prediction strategies of a model evaluated on inputs that require gradients fails
         # with RuntimeError (Only Tensors created explicitly by the user (graph leaves) support the deepcopy
         # protocol at the moment). Overwriting this method make sure that the prediction strategies of a
         # model are set to None upon deepcopying.
@@ -80,7 +80,7 @@ class DefaultPredictionStrategy(object):
         # where S S^T = (K_XX + sigma^2 I)^-1
         return test_train_covar.matmul(precomputed_cache)
 
-    def get_fantasy_strategy(self, inputs, targets, full_inputs, full_targets, full_output):
+    def get_fantasy_strategy(self, inputs, targets, full_inputs, full_targets, full_output, **kwargs):
         """
         Returns a new PredictionStrategy that incorporates the specified inputs and targets as new training data.
 
@@ -108,9 +108,11 @@ class DefaultPredictionStrategy(object):
         # Evaluate fant x train and fant x fant covariance matrices, leave train x train unevaluated.
         fant_fant_covar = full_covar[..., num_train:, num_train:]
         fant_mean = full_mean[..., num_train:]
-        mvn = self.likelihood(self.train_prior_dist.__class__(fant_mean, fant_fant_covar), inputs)
-        fant_fant_covar = mvn.covariance_matrix
+        mvn = self.train_prior_dist.__class__(fant_mean, fant_fant_covar)
+        self.likelihood = self.likelihood.get_fantasy_likelihood(**kwargs)
+        mvn_obs = self.likelihood(mvn, inputs, **kwargs)
 
+        fant_fant_covar = mvn_obs.covariance_matrix
         fant_train_covar = delazify(full_covar[..., num_train:, :num_train])
 
         self.fantasy_inputs = inputs
@@ -134,10 +136,11 @@ class DefaultPredictionStrategy(object):
         # Solve for "b", the lower portion of the *new* \\alpha corresponding to the fantasy points.
         schur_complement = fant_fant_covar - fant_train_covar.matmul(fant_solve)
         small_system_rhs = targets - fant_mean - fant_train_covar.matmul(self.mean_cache)
+        small_system_rhs = small_system_rhs.unsqueeze(-1)
         # Schur complement of a spd matrix is guaranteed to be positive definite
         if small_system_rhs.requires_grad or schur_complement.requires_grad:
             # TODO: Delete this part of the if statement when PyTorch implements cholesky_solve derivative.
-            fant_cache_lower = torch.gesv(small_system_rhs.unsqueeze(-1), schur_complement)[0]
+            fant_cache_lower = torch.gesv(small_system_rhs, schur_complement)[0]
         else:
             fant_cache_lower = cholesky_solve(small_system_rhs, torch.cholesky(schur_complement))
 
@@ -355,7 +358,7 @@ class InterpolatedPredictionStrategy(DefaultPredictionStrategy):
         res = left_interp(test_interp_indices, test_interp_values, precomputed_cache)
         return res
 
-    def get_fantasy_strategy(self, inputs, targets, full_inputs, full_targets, full_output):
+    def get_fantasy_strategy(self, inputs, targets, full_inputs, full_targets, full_output, **kwargs):
         raise NotImplementedError(
             "Fantasy observation updates not yet supported for models using InterpolatedLazyTensors"
         )


Test patch:
diff --git a/test/examples/test_fixed_noise_fanatasy_updates.py b/test/examples/test_fixed_noise_fanatasy_updates.py
new file mode 100644
index 00000000..e2beb16a
--- /dev/null
+++ b/test/examples/test_fixed_noise_fanatasy_updates.py
@@ -0,0 +1,174 @@
+#!/usr/bin/env python3
+
+import unittest
+from math import exp, pi
+from test._utils import least_used_cuda_device
+
+import gpytorch
+import torch
+from gpytorch.distributions import MultivariateNormal
+from gpytorch.kernels import RBFKernel, ScaleKernel
+from gpytorch.likelihoods import FixedNoiseGaussianLikelihood
+from gpytorch.likelihoods.noise_models import FixedGaussianNoise
+from gpytorch.means import ConstantMean
+from gpytorch.priors import SmoothedBoxPrior
+from torch import optim
+from test._base_test_case import BaseTestCase
+
+
+class ExactGPModel(gpytorch.models.ExactGP):
+    def __init__(self, train_inputs, train_targets, likelihood):
+        super(ExactGPModel, self).__init__(train_inputs, train_targets, likelihood)
+        self.mean_module = ConstantMean(prior=SmoothedBoxPrior(-1, 1))
+        self.covar_module = ScaleKernel(RBFKernel(lengthscale_prior=SmoothedBoxPrior(exp(-3), exp(3), sigma=0.1)))
+
+    def forward(self, x):
+        mean_x = self.mean_module(x)
+        covar_x = self.covar_module(x)
+        return MultivariateNormal(mean_x, covar_x)
+
+
+class TestFixedNoiseFantasies(BaseTestCase, unittest.TestCase):
+    seed = 1
+
+    def _get_data(self, cuda=False, num_data=11, add_noise=False):
+        device = torch.device("cuda") if cuda else torch.device("cpu")
+        # Simple training data: let's try to learn a sine function
+        train_x = torch.linspace(0, 1, num_data, device=device)
+        train_y = torch.sin(train_x * (2 * pi))
+        if add_noise:
+            train_y.add_(torch.randn_like(train_x).mul_(0.1))
+        test_x = torch.linspace(0, 1, 51, device=device)
+        test_y = torch.sin(test_x * (2 * pi))
+        return train_x, test_x, train_y, test_y
+
+    def test_fixed_noise_fantasy_updates_cuda(self):
+        if torch.cuda.is_available():
+            with least_used_cuda_device():
+                self.test_fantasy_updates(cuda=True)
+
+    def test_fixed_noise_fantasy_updates(self, cuda=False):
+        train_x, test_x, train_y, test_y = self._get_data(cuda=cuda)
+        noise = torch.full_like(train_y, 2e-4)
+        test_noise = torch.full_like(test_y, 3e-4)
+
+        likelihood = FixedNoiseGaussianLikelihood(noise)
+        gp_model = ExactGPModel(train_x, train_y, likelihood)
+        mll = gpytorch.ExactMarginalLogLikelihood(likelihood, gp_model)
+        gp_model.covar_module.base_kernel.initialize(lengthscale=exp(1))
+        gp_model.mean_module.initialize(constant=0)
+
+        if cuda:
+            gp_model.cuda()
+            likelihood.cuda()
+
+        # Find optimal model hyperparameters
+        gp_model.train()
+        likelihood.train()
+        optimizer = optim.Adam(list(gp_model.parameters()) + list(likelihood.parameters()), lr=0.15)
+        for _ in range(50):
+            optimizer.zero_grad()
+            with gpytorch.settings.debug(False):
+                output = gp_model(train_x)
+            loss = -mll(output, train_y)
+            loss.backward()
+            optimizer.step()
+
+        for param in gp_model.parameters():
+            self.assertTrue(param.grad is not None)
+            self.assertGreater(param.grad.norm().item(), 0)
+        optimizer.step()
+
+        train_x.requires_grad = True
+        gp_model.set_train_data(train_x, train_y)
+        with gpytorch.settings.fast_pred_var(), gpytorch.settings.detach_test_caches(False):
+            # Test the model
+            gp_model.eval()
+            likelihood.eval()
+            test_function_predictions = likelihood(gp_model(test_x), noise=test_noise)
+            test_function_predictions.mean.sum().backward()
+
+            real_fant_x_grad = train_x.grad[5:].clone()
+            train_x.grad = None
+            train_x.requires_grad = False
+            gp_model.set_train_data(train_x, train_y)
+
+            # Cut data down, and then add back via the fantasy interface
+            gp_model.set_train_data(train_x[:5], train_y[:5], strict=False)
+            gp_model.likelihood.noise_covar = FixedGaussianNoise(noise=noise[:5])
+            likelihood(gp_model(test_x), noise=test_noise)
+
+            fantasy_x = train_x[5:].clone().detach().requires_grad_(True)
+            fant_model = gp_model.get_fantasy_model(fantasy_x, train_y[5:], noise=noise[5:])
+
+            fant_function_predictions = likelihood(fant_model(test_x), noise=test_noise)
+
+            self.assertAllClose(test_function_predictions.mean, fant_function_predictions.mean, atol=1e-4)
+
+            fant_function_predictions.mean.sum().backward()
+            self.assertTrue(fantasy_x.grad is not None)
+
+            relative_error = torch.norm(real_fant_x_grad - fantasy_x.grad) / fantasy_x.grad.norm()
+            self.assertLess(relative_error, 15e-1)  # This was only passing by a hair before
+
+    def test_fixed_noise_fantasy_updates_batch_cuda(self):
+        if torch.cuda.is_available():
+            with least_used_cuda_device():
+                self.test_fantasy_updates_batch(cuda=True)
+
+    def test_fixed_noise_fantasy_updates_batch(self, cuda=False):
+        train_x, test_x, train_y, test_y = self._get_data(cuda=cuda)
+        noise = torch.full_like(train_y, 2e-4)
+        test_noise = torch.full_like(test_y, 3e-4)
+
+        likelihood = FixedNoiseGaussianLikelihood(noise)
+        gp_model = ExactGPModel(train_x, train_y, likelihood)
+        mll = gpytorch.ExactMarginalLogLikelihood(likelihood, gp_model)
+        gp_model.covar_module.base_kernel.initialize(lengthscale=exp(1))
+        gp_model.mean_module.initialize(constant=0)
+
+        if cuda:
+            gp_model.cuda()
+            likelihood.cuda()
+
+        # Find optimal model hyperparameters
+        gp_model.train()
+        likelihood.train()
+        optimizer = optim.Adam(list(gp_model.parameters()) + list(likelihood.parameters()), lr=0.15)
+        for _ in range(50):
+            optimizer.zero_grad()
+            with gpytorch.settings.debug(False):
+                output = gp_model(train_x)
+            loss = -mll(output, train_y)
+            loss.backward()
+            optimizer.step()
+
+        for param in gp_model.parameters():
+            self.assertTrue(param.grad is not None)
+            self.assertGreater(param.grad.norm().item(), 0)
+        optimizer.step()
+
+        with gpytorch.settings.fast_pred_var():
+            # Test the model
+            gp_model.eval()
+            likelihood.eval()
+            test_function_predictions = likelihood(gp_model(test_x), noise=test_noise)
+
+            # Cut data down, and then add back via the fantasy interface
+            gp_model.set_train_data(train_x[:5], train_y[:5], strict=False)
+            gp_model.likelihood.noise_covar = FixedGaussianNoise(noise=noise[:5])
+            likelihood(gp_model(test_x), noise=test_noise)
+
+            fantasy_x = train_x[5:].clone().unsqueeze(0).unsqueeze(-1).repeat(3, 1, 1).requires_grad_(True)
+            fantasy_y = train_y[5:].unsqueeze(0).repeat(3, 1)
+            fant_model = gp_model.get_fantasy_model(fantasy_x, fantasy_y, noise=noise[5:].unsqueeze(0).repeat(3, 1))
+            fant_function_predictions = likelihood(fant_model(test_x), noise=test_noise)
+
+            self.assertAllClose(test_function_predictions.mean, fant_function_predictions.mean[0], atol=1e-4)
+
+            fant_function_predictions.mean.sum().backward()
+            self.assertTrue(fantasy_x.grad is not None)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/test/likelihoods/test_gaussian_likelihood.py b/test/likelihoods/test_gaussian_likelihood.py
index 22439e1f..31078af4 100644
--- a/test/likelihoods/test_gaussian_likelihood.py
+++ b/test/likelihoods/test_gaussian_likelihood.py
@@ -69,7 +69,7 @@ class TestFixedNoiseGaussianLikelihood(BaseLikelihoodTestCase, unittest.TestCase
                 lkhd(mvn)
             # test __call__ w/ observation noise
             obs_noise = 0.1 + torch.rand(5, device=device, dtype=dtype)
-            out = lkhd(mvn, observation_noise=obs_noise)
+            out = lkhd(mvn, noise=obs_noise)
             self.assertTrue(torch.allclose(out.variance, 1 + obs_noise))