Repository: cornellius-gp/gpytorch

Issue:
[Bug]
# 🐛 Bug

Hello GPyTorch team,

I am currently working on a project that requires using Gaussian processes to handle multi-task regression. I am using your ApproximateGP model with IndependentMultitaskVariationalStrategy, as shown in your [SVGP Multitask GP Regression example](https://docs.gpytorch.ai/en/latest/examples/04_Variational_and_Approximate_GPs/SVGP_Multitask_GP_Regression.html). However, I've run into an issue when trying to update the model with new training data.

My use-case necessitates setting new training data for the Gaussian process to condition on, without refitting the model parameters. I understand that ExactGP models have the set_train_data method which suits this purpose perfectly, but ApproximateGP does not seem to have a similar functionality. To work around this, I attempted to use the get_fantasy_model method, hoping to generate a new model conditioned on both the existing and new data.

Unfortunately, when I attempted to use get_fantasy_model with my ApproximateGP model (initialized with IndependentMultitaskVariationalStrategy), I encountered an AttributeError indicating that the IndependentMultitaskVariationalStrategy object does not have an attribute __name__. It appears that get_fantasy_model is not fully compatible with IndependentMultitaskVariationalStrategy, or perhaps with ApproximateGP models in general.

Given this situation, I'd appreciate if you could provide some guidance or potential workarounds. Furthermore, would you consider adding a functionality similar to set_train_data for ApproximateGP models, or possibly extending the compatibility of get_fantasy_model to handle ApproximateGP and IndependentMultitaskVariationalStrategy? I believe these additions would make GPyTorch more flexible in handling a wider range of use-cases involving variational and multitask GPs.

Thanks for your time and your continuous work on this great library.

Best regards,
İlter Onat Korkmaz
## To reproduce

** Code snippet to reproduce **
```python
import torch
import gpytorch

class IndependentMultitaskGPModel(gpytorch.models.ApproximateGP):
    def __init__(self, num_tasks):
        inducing_points = torch.rand(num_tasks, 16, 1)
        variational_distribution = gpytorch.variational.CholeskyVariationalDistribution(
            inducing_points.size(-2), batch_shape=torch.Size([num_tasks])
        )
        variational_strategy = gpytorch.variational.IndependentMultitaskVariationalStrategy(
            gpytorch.variational.VariationalStrategy(
                self, inducing_points, variational_distribution, learn_inducing_locations=True
            ),
            num_tasks=num_tasks,
        )
        super().__init__(variational_strategy)
        self.mean_module = gpytorch.means.ConstantMean(batch_shape=torch.Size([num_tasks]))
        self.covar_module = gpytorch.kernels.ScaleKernel(
            gpytorch.kernels.RBFKernel(batch_shape=torch.Size([num_tasks])),
            batch_shape=torch.Size([num_tasks])
        )

    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)

num_tasks = 4
model = IndependentMultitaskGPModel(num_tasks)
x_new = torch.rand(num_tasks, 10, 1)
y_new = torch.rand(num_tasks, 10)
fantasy_model = model.get_fantasy_model(x_new, y_new)  # Raises AttributeError

```

** Stack trace/error message **
```
---------------------------------------------------------------------------

Output is truncated. View as a scrollable element or open in a text editor. Adjust cell output settings...
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
File ~/env/lib/python3.10/site-packages/gpytorch/module.py:437, in Module.__getattr__(self, name)
    436 try:
--> 437     return super().__getattribute__(name)
    438 except AttributeError:

AttributeError: 'IndependentMultitaskVariationalStrategy' object has no attribute '__name__'

During handling of the above exception, another exception occurred:

AttributeError                            Traceback (most recent call last)
Cell In[1], line 32
     30 x_new = torch.rand(num_tasks, 10, 1)
     31 y_new = torch.rand(num_tasks, 10)
---> 32 fantasy_model = model.get_fantasy_model(x_new, y_new)  # Raises AttributeError

File ~/env/lib/python3.10/site-packages/gpytorch/models/approximate_gp.py:103, in ApproximateGP.get_fantasy_model(self, inputs, targets, **kwargs)
     78 def get_fantasy_model(self, inputs, targets, **kwargs):
     79     r"""
     80     Returns a new GP model that incorporates the specified inputs and targets as new training data using
     81     online variational conditioning (OVC).
   (...)
    101 
    102     """
...
   1613         return modules[name]
-> 1614 raise AttributeError("'{}' object has no attribute '{}'".format(
   1615     type(self).__name__, name))

AttributeError: 'IndependentMultitaskVariationalStrategy' object has no attribute '__name__'


```

## Expected Behavior

The code returns an ExactGP model as the fantasy model.

## System information

**Please complete the following information:**
- <!-- GPyTorch Version (run `print(gpytorch.__version__)` --> 1.10
- <!-- PyTorch Version (run `print(torch.__version__)` --> 2.0.0+cu117
- <!-- Computer OS --> Windows 11

## Additional context
None









































































































Discussion:
There was a small typo in our error handling (I'm putting up a PR now). This is the error that you should have seen:

```
NotImplementedError: No fantasy model support for IndependentMultitaskVariationalStrategy. Only VariationalStrategy and UnwhitenedVariationalStrategy are currently supported.
````

> My use-case necessitates setting new training data for the Gaussian process to condition on, without refitting the model parameters. I understand that ExactGP models have the set_train_data method which suits this purpose perfectly, but ApproximateGP does not seem to have a similar functionality. To work around this, I attempted to use the get_fantasy_model method, hoping to generate a new model conditioned on both the existing and new data.

Currently we don't support this for IndependentMultitaskVariationalStrategy, though in theory it shouldn't be too hard to set up (it would essentially just require calling through to the base variational strategy). We'd accept a PR if you'd be up to implementing it!

Resolving patch:
diff --git a/gpytorch/variational/_variational_strategy.py b/gpytorch/variational/_variational_strategy.py
--- a/gpytorch/variational/_variational_strategy.py
+++ b/gpytorch/variational/_variational_strategy.py
@@ -264,15 +264,21 @@ def get_fantasy_model(
         # whitened / unwhitened variational strategies
         if not self.has_fantasy_strategy:
             raise NotImplementedError(
-                "No fantasy model support for ",
-                self.__name__,
-                ". Only VariationalStrategy and UnwhitenedVariationalStrategy are currently supported.",
+                f"No fantasy model support for {self.__class__.__name__}. "
+                "Only VariationalStrategy and UnwhitenedVariationalStrategy are currently supported."
             )
+        else:
+            from . import CholeskyVariationalDistribution  # Circular import otherwise
+
+            if not isinstance(self._variational_distribution, CholeskyVariationalDistribution):
+                raise NotImplementedError(
+                    "Fantasy models are only support for variational models with CholeskyVariationalDistribution."
+                )
+
         if not isinstance(self.model.likelihood, GaussianLikelihood):
             raise NotImplementedError(
-                "No fantasy model support for ",
-                self.model.likelihood,
-                ". Only GaussianLikelihoods are currently supported.",
+                f"No fantasy model support for {self.model.likelihood.__class__.__name__}. "
+                "Only GaussianLikelihoods are currently supported."
             )
         # we assume that either the user has given the model a mean_module and a covar_module
         # or that it will be passed into the get_fantasy_model function. we check for these.
@@ -280,16 +286,16 @@ def get_fantasy_model(
             mean_module = getattr(self.model, "mean_module", None)
             if mean_module is None:
                 raise ModuleNotFoundError(
-                    "Either you must provide a mean_module as input to get_fantasy_model",
-                    "or it must be an attribute of the model called mean_module.",
+                    "Either you must provide a mean_module as input to get_fantasy_model "
+                    "or it must be an attribute of the model called mean_module."
                 )
         if covar_module is None:
             covar_module = getattr(self.model, "covar_module", None)
             if covar_module is None:
                 # raise an error
                 raise ModuleNotFoundError(
-                    "Either you must provide a covar_module as input to get_fantasy_model",
-                    "or it must be an attribute of the model called covar_module.",
+                    "Either you must provide a covar_module as input to get_fantasy_model "
+                    "or it must be an attribute of the model called covar_module."
                 )
 
         # first we construct an exact model over the inducing points with the inducing covariance
diff --git a/gpytorch/variational/nearest_neighbor_variational_strategy.py b/gpytorch/variational/nearest_neighbor_variational_strategy.py
--- a/gpytorch/variational/nearest_neighbor_variational_strategy.py
+++ b/gpytorch/variational/nearest_neighbor_variational_strategy.py
@@ -9,7 +9,8 @@
 from torch import LongTensor, Tensor
 
 from ..distributions import MultivariateNormal
-from ..models import ApproximateGP
+from ..models import ApproximateGP, ExactGP
+from ..module import Module
 from ..utils.errors import CachingError
 from ..utils.memoize import add_to_cache, cached, pop_from_cache
 from ..utils.nearest_neighbors import NNUtil
@@ -238,6 +239,19 @@ def forward(
             # Return the distribution
             return MultivariateNormal(predictive_mean, DiagLinearOperator(predictive_var))
 
+    def get_fantasy_model(
+        self,
+        inputs: Tensor,
+        targets: Tensor,
+        mean_module: Optional[Module] = None,
+        covar_module: Optional[Module] = None,
+        **kwargs,
+    ) -> ExactGP:
+        raise NotImplementedError(
+            f"No fantasy model support for {self.__class__.__name__}. "
+            "Only VariationalStrategy and UnwhitenedVariationalStrategy are currently supported."
+        )
+
     def _set_training_iterator(self) -> None:
         self._training_indices_iter = 0
         training_indices = torch.randperm(self.M - self.k, device=self.inducing_points.device) + self.k


Test patch:
diff --git a/test/variational/test_ciq_variational_strategy.py b/test/variational/test_ciq_variational_strategy.py
--- a/test/variational/test_ciq_variational_strategy.py
+++ b/test/variational/test_ciq_variational_strategy.py
@@ -38,7 +38,7 @@ def test_eval_iteration(self, *args, **kwargs):
         self.assertEqual(ciq_mock.call_count, 2)  # One for each evaluation call
 
     def test_fantasy_call(self, *args, **kwargs):
-        with self.assertRaises(AttributeError):
+        with self.assertRaises(NotImplementedError):
             super().test_fantasy_call(*args, **kwargs)
 
 
diff --git a/test/variational/test_grid_interpolation_variational_strategy.py b/test/variational/test_grid_interpolation_variational_strategy.py
--- a/test/variational/test_grid_interpolation_variational_strategy.py
+++ b/test/variational/test_grid_interpolation_variational_strategy.py
@@ -76,7 +76,7 @@ def test_eval_iteration(self, *args, **kwargs):
         self.assertFalse(ciq_mock.called)
 
     def test_fantasy_call(self, *args, **kwargs):
-        with self.assertRaises(AttributeError):
+        with self.assertRaises(NotImplementedError):
             super().test_fantasy_call(*args, **kwargs)
 
 
diff --git a/test/variational/test_independent_multitask_variational_strategy.py b/test/variational/test_independent_multitask_variational_strategy.py
--- a/test/variational/test_independent_multitask_variational_strategy.py
+++ b/test/variational/test_independent_multitask_variational_strategy.py
@@ -61,7 +61,7 @@ def test_eval_iteration(self, *args, expected_batch_shape=None, **kwargs):
         super().test_eval_iteration(*args, expected_batch_shape=expected_batch_shape, **kwargs)
 
     def test_fantasy_call(self, *args, **kwargs):
-        with self.assertRaises(AttributeError):
+        with self.assertRaises(NotImplementedError):
             super().test_fantasy_call(*args, **kwargs)
 
 
diff --git a/test/variational/test_lmc_variational_strategy.py b/test/variational/test_lmc_variational_strategy.py
--- a/test/variational/test_lmc_variational_strategy.py
+++ b/test/variational/test_lmc_variational_strategy.py
@@ -69,7 +69,7 @@ def test_eval_iteration(self, *args, expected_batch_shape=None, **kwargs):
         self.assertFalse(ciq_mock.called)
 
     def test_fantasy_call(self, *args, **kwargs):
-        with self.assertRaises(AttributeError):
+        with self.assertRaises(NotImplementedError):
             super().test_fantasy_call(*args, **kwargs)
 
 
diff --git a/test/variational/test_nearest_neighbor_variational_strategy.py b/test/variational/test_nearest_neighbor_variational_strategy.py
--- a/test/variational/test_nearest_neighbor_variational_strategy.py
+++ b/test/variational/test_nearest_neighbor_variational_strategy.py
@@ -42,15 +42,13 @@ def _make_model_and_likelihood(
         distribution_cls=gpytorch.variational.MeanFieldVariationalDistribution,
         constant_mean=True,
     ):
-        class _VNNGPRegressionModel(gpytorch.models.GP):
+        class _VNNGPRegressionModel(gpytorch.models.ApproximateGP):
             def __init__(self, inducing_points, k, training_batch_size):
-                super(_VNNGPRegressionModel, self).__init__()
-
                 variational_distribution = distribution_cls(num_inducing, batch_shape=batch_shape)
-
-                self.variational_strategy = strategy_cls(
+                variational_strategy = strategy_cls(
                     self, inducing_points, variational_distribution, k=k, training_batch_size=training_batch_size
                 )
+                super().__init__(variational_strategy)
 
                 if constant_mean:
                     self.mean_module = gpytorch.means.ConstantMean()
@@ -252,7 +250,7 @@ def test_eval_larger_pred_batch(self):
         )
 
     def test_fantasy_call(self, *args, **kwargs):
-        with self.assertRaises(AttributeError):
+        with self.assertRaises(NotImplementedError):
             super().test_fantasy_call(*args, **kwargs)
 
 
diff --git a/test/variational/test_orthogonally_decoupled_variational_strategy.py b/test/variational/test_orthogonally_decoupled_variational_strategy.py
--- a/test/variational/test_orthogonally_decoupled_variational_strategy.py
+++ b/test/variational/test_orthogonally_decoupled_variational_strategy.py
@@ -58,7 +58,7 @@ def test_eval_iteration(self, *args, **kwargs):
         self.assertEqual(cholesky_mock.call_count, 1)  # One to compute cache, that's it!
 
     def test_fantasy_call(self, *args, **kwargs):
-        with self.assertRaises(AttributeError):
+        with self.assertRaises(NotImplementedError):
             super().test_fantasy_call(*args, **kwargs)
 
 
diff --git a/test/variational/test_unwhitened_variational_strategy.py b/test/variational/test_unwhitened_variational_strategy.py
--- a/test/variational/test_unwhitened_variational_strategy.py
+++ b/test/variational/test_unwhitened_variational_strategy.py
@@ -45,7 +45,7 @@ def test_fantasy_call(self, *args, **kwargs):
         if self.distribution_cls is gpytorch.variational.CholeskyVariationalDistribution:
             return super().test_fantasy_call(*args, **kwargs)
 
-        with self.assertRaises(AttributeError):
+        with self.assertRaises(NotImplementedError):
             super().test_fantasy_call(*args, **kwargs)
 
 
diff --git a/test/variational/test_variational_strategy.py b/test/variational/test_variational_strategy.py
--- a/test/variational/test_variational_strategy.py
+++ b/test/variational/test_variational_strategy.py
@@ -42,7 +42,7 @@ def test_fantasy_call(self, *args, **kwargs):
         if self.distribution_cls is gpytorch.variational.CholeskyVariationalDistribution:
             return super().test_fantasy_call(*args, **kwargs)
 
-        with self.assertRaises(AttributeError):
+        with self.assertRaises(NotImplementedError):
             super().test_fantasy_call(*args, **kwargs)