Support SaasFullyBayesianSingleTaskGP in prune_inferior_points (meta-pytorch#1260)

Summary:
Pull Request resolved: meta-pytorch#1260

See title

Reviewed By: saitcakmak

Differential Revision: D37091706

fbshipit-source-id: c091683cdb52a5d2f21f406d98589f94d9fa218b

Author: dme65

botorch/acquisition/utils.py

@@ -26,12 +26,14 @@
 )
 from botorch.exceptions.errors import UnsupportedError
 from botorch.exceptions.warnings import SamplingWarning
+from botorch.models.fully_bayesian import MCMC_DIM
 from botorch.models.model import Model
 from botorch.sampling.samplers import IIDNormalSampler, MCSampler, SobolQMCNormalSampler
 from botorch.utils.multi_objective.box_decompositions.non_dominated import (
     FastNondominatedPartitioning,
     NondominatedPartitioning,
 )
+from botorch.utils.transforms import is_fully_bayesian
 from torch import Tensor
 from torch.quasirandom import SobolEngine
 
@@ -320,6 +322,10 @@ def prune_inferior_points(
         with `N_nz` the number of points in `X` that have non-zero (empirical,
         under `num_samples` samples) probability of being the best point.
     """
+    if marginalize_dim is None and is_fully_bayesian(model):
+        # TODO: Properly deal with marginalizing fully Bayesian models
+        marginalize_dim = MCMC_DIM
+
     if X.ndim > 2:
         # TODO: support batched inputs (req. dealing with ragged tensors)
         raise UnsupportedError(

botorch/utils/transforms.py

@@ -183,12 +183,12 @@ def is_fully_bayesian(model: Model) -> bool:
     Returns:
         True if at least one model is a `SaasFullyBayesianSingleTaskGP`
     """
-    from botorch.models import ModelList
+    from botorch.models import ModelList, ModelListGP
     from botorch.models.fully_bayesian import SaasFullyBayesianSingleTaskGP
 
     if isinstance(model, SaasFullyBayesianSingleTaskGP):
         return True
-    elif isinstance(model, ModelList) and any(
+    elif isinstance(model, (ModelList, ModelListGP)) and any(
         isinstance(m, SaasFullyBayesianSingleTaskGP) for m in model.models
     ):
         return True

test/models/test_fully_bayesian.py

@@ -27,6 +27,7 @@
     qExpectedHypervolumeImprovement,
     qNoisyExpectedHypervolumeImprovement,
 )
+from botorch.acquisition.utils import prune_inferior_points
 from botorch.models import ModelList, ModelListGP
 from botorch.models.deterministic import GenericDeterministicModel
 from botorch.models.fully_bayesian import (
@@ -423,6 +424,10 @@ def test_acquisition_functions(self):
                 test_X = torch.rand(*batch_shape, 1, 4, **tkwargs)
                 self.assertEqual(acqf(test_X).shape, torch.Size(batch_shape))
 
+        # Test prune_inferior_points
+        X_pruned = prune_inferior_points(model=model, X=train_X)
+        self.assertTrue(X_pruned.ndim == 2 and X_pruned.shape[-1] == 4)
+
         # Test prune_inferior_points_multi_objective
         for model_list in [ModelListGP(model, model), ModelList(deterministic, model)]:
             X_pruned = prune_inferior_points_multi_objective(

test/utils/test_transforms.py

@@ -8,11 +8,19 @@
 from typing import Any
 
 import torch
+from botorch.models import (
+    GenericDeterministicModel,
+    ModelList,
+    ModelListGP,
+    SaasFullyBayesianSingleTaskGP,
+    SingleTaskGP,
+)
 from botorch.models.model import Model
 from botorch.utils.testing import BotorchTestCase, MockModel, MockPosterior
 from botorch.utils.transforms import (
     _verify_output_shape,
     concatenate_pending_points,
+    is_fully_bayesian,
     match_batch_shape,
     normalize,
     normalize_indices,
@@ -278,3 +286,23 @@ def test_squeeze_last_dim(self):
             Y_squeezed = squeeze_last_dim(Y=Y)
             self.assertTrue(any(issubclass(w.category, DeprecationWarning) for w in ws))
         self.assertTrue(torch.equal(Y_squeezed, Y.squeeze(-1)))
+
+
+class TestIsFullyBayesian(BotorchTestCase):
+    def test_is_fully_bayesian(self):
+        X, Y = torch.rand(3, 2), torch.randn(3, 1)
+        saas = SaasFullyBayesianSingleTaskGP(train_X=X, train_Y=Y)
+        vanilla_gp = SingleTaskGP(train_X=X, train_Y=Y)
+        deterministic = GenericDeterministicModel(f=lambda x: x)
+        # Single model
+        self.assertTrue(is_fully_bayesian(model=saas))
+        self.assertFalse(is_fully_bayesian(model=vanilla_gp))
+        self.assertFalse(is_fully_bayesian(model=deterministic))
+        # ModelListGP
+        self.assertTrue(is_fully_bayesian(model=ModelListGP(saas, saas)))
+        self.assertTrue(is_fully_bayesian(model=ModelListGP(saas, vanilla_gp)))
+        self.assertFalse(is_fully_bayesian(model=ModelListGP(vanilla_gp, vanilla_gp)))
+        # ModelList
+        self.assertTrue(is_fully_bayesian(model=ModelList(saas, saas)))
+        self.assertTrue(is_fully_bayesian(model=ModelList(saas, deterministic)))
+        self.assertFalse(is_fully_bayesian(model=ModelList(vanilla_gp, deterministic)))