EnsembleMapSaasGP -> EnsembleMapSaasSingleTaskGP (meta-pytorch#3038)

Summary:
Pull Request resolved: meta-pytorch#3038

Realized a little bit too late that `EnsembleMapSaasGP` will be confusing once this model is extended to an MTGP. Let's rename it consistently with our other models before it's too late :)

Reviewed By: Balandat

Differential Revision: D83796108

fbshipit-source-id: 9618f63fcec7117af4afa324379a96c8b49790f2

Author: David Eriksson

botorch/models/__init__.py

@@ -23,7 +23,7 @@
 from botorch.models.map_saas import (
     add_saas_prior,
     AdditiveMapSaasSingleTaskGP,
-    EnsembleMapSaasGP,
+    EnsembleMapSaasSingleTaskGP,
 )
 from botorch.models.model import ModelList
 from botorch.models.model_list_gp_regression import ModelListGP
@@ -36,9 +36,7 @@
     "AffineDeterministicModel",
     "AffineFidelityCostModel",
     "ApproximateGPyTorchModel",
-    "EnsembleMapSaasGP",
-    "SaasFullyBayesianSingleTaskGP",
-    "SaasFullyBayesianMultiTaskGP",
+    "EnsembleMapSaasSingleTaskGP",
     "GenericDeterministicModel",
     "HigherOrderGP",
     "KroneckerMultiTaskGP",
@@ -49,6 +47,8 @@
     "PairwiseGP",
     "PairwiseLaplaceMarginalLogLikelihood",
     "PosteriorMeanModel",
+    "SaasFullyBayesianMultiTaskGP",
+    "SaasFullyBayesianSingleTaskGP",
     "SingleTaskGP",
     "SingleTaskMultiFidelityGP",
     "SingleTaskVariationalGP",

botorch/models/map_saas.py

@@ -39,9 +39,9 @@ def __init__(self, tau: Tensor | float | None = None):
         Args:
             tau: Value of the global shrinkage parameter. If `None`, the tau will be
                 a free parameter and inferred from the data.
-                Tau can be a tensor for batched models, like `EnsembleMapSaasGP`,
-                where each batch has a different sparsity prior. If tau is a tensor,
-                it must have shape `batch_shape`.
+                Tau can be a tensor for batched models, like
+                `EnsembleMapSaasSingleTaskGP`, where each batch has a different
+                sparsity prior. If tau is a tensor, it must have shape `batch_shape`.
         """
         self._tau = torch.as_tensor(tau) if tau is not None else None
 
@@ -427,7 +427,7 @@ def __init__(
         self.to(dtype=train_X.dtype, device=train_X.device)
 
 
-class EnsembleMapSaasGP(SingleTaskGP):
+class EnsembleMapSaasSingleTaskGP(SingleTaskGP):
     _is_ensemble = True
 
     def __init__(
@@ -440,9 +440,9 @@ def __init__(
         outcome_transform: OutcomeTransform | _DefaultType | None = DEFAULT,
         input_transform: InputTransform | None = None,
     ) -> None:
-        """Instantiates an ``EnsembleMapSaasGP``, which is a batched ensemble of
-        ``SingleTaskGP``s with the Matern-5/2 kernel and a SAAS prior. The model is
-        intended to be trained with ``ExactMarginalLogLikelihood`` and
+        """Instantiates an ``EnsembleMapSaasSingleTaskGP``, which is a batched
+        ensemble of ``SingleTaskGP``s with the Matern-5/2 kernel and a SAAS prior.
+        The model is intended to be trained with ``ExactMarginalLogLikelihood`` and
         ``fit_gpytorch_mll``. Under the hood, the model is equivalent to a
         multi-output ``BatchedMultiOutputGPyTorchModel``, but it produces a
         ``MixtureGaussiaPosterior``, which leads to ensembling of the model outputs.

test/models/test_map_saas.py

@@ -22,7 +22,7 @@
 from botorch.models.map_saas import (
     add_saas_prior,
     AdditiveMapSaasSingleTaskGP,
-    EnsembleMapSaasGP,
+    EnsembleMapSaasSingleTaskGP,
     get_additive_map_saas_covar_module,
     get_gaussian_likelihood_with_gamma_prior,
     get_mean_module_with_normal_prior,
@@ -527,7 +527,7 @@ def test_emsemble_map_saas(self) -> None:
                 }
             else:
                 extra_inputs = {}
-            model = EnsembleMapSaasGP(
+            model = EnsembleMapSaasSingleTaskGP(
                 train_X=train_X, train_Y=train_Y, num_taus=num_taus, **extra_inputs
             )
             sample_all_priors(model)  # Checks that the prior is configured correctly.
@@ -553,16 +553,20 @@ def test_emsemble_map_saas(self) -> None:
 
     def test_ensemble_map_saas_validation(self) -> None:
         with self.assertRaisesRegex(ValueError, "Expected taus to be of shape"):
-            EnsembleMapSaasGP(
+            EnsembleMapSaasSingleTaskGP(
                 train_X=torch.rand(5, 3),
                 train_Y=torch.rand(5, 1),
                 num_taus=3,
                 taus=torch.rand(2),
             )
         with self.assertRaisesRegex(UnsupportedError, "only supports single-output"):
-            EnsembleMapSaasGP(train_X=torch.rand(5, 3), train_Y=torch.rand(5, 2))
+            EnsembleMapSaasSingleTaskGP(
+                train_X=torch.rand(5, 3), train_Y=torch.rand(5, 2)
+            )
         with self.assertRaisesRegex(UnsupportedError, "only supports 2D inputs"):
-            EnsembleMapSaasGP(train_X=torch.rand(2, 5, 3), train_Y=torch.rand(2, 5, 1))
+            EnsembleMapSaasSingleTaskGP(
+                train_X=torch.rand(2, 5, 3), train_Y=torch.rand(2, 5, 1)
+            )
 
 
 class TestAdditiveMapSaasSingleTaskGP(BotorchTestCase):