Fix error messages for ApproximateGP.get_fantasy_model (#2374)

[Fixes #2370]

Author: gpleiss

gpytorch/variational/_variational_strategy.py

@@ -264,32 +264,38 @@ 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.
         if mean_module is None:
             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

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/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)
 
 

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)
 
 

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)
 
 

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)
 
 

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)
 
 

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)
 
 

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)
 
 

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)