Add construct_inputs to EnsembleMapSaasSingleTaskGP (meta-pytorch#3040)

Summary:
Pull Request resolved: meta-pytorch#3040

We need this to pass in `num_taus` from MBM.

Reviewed By: Balandat

Differential Revision: D83442381

fbshipit-source-id: 6a83e2cd79d38efcb3605a34608cc00f1a222902

Author: David Eriksson

botorch/models/map_saas.py

@@ -16,6 +16,8 @@
 )
 from botorch.posteriors.fully_bayesian import GaussianMixturePosterior, MCMC_DIM
 from botorch.utils.constraints import LogTransformedInterval
+from botorch.utils.containers import BotorchContainer
+from botorch.utils.datasets import SupervisedDataset
 from botorch.utils.types import _DefaultType, DEFAULT
 from gpytorch.constraints import Interval
 from gpytorch.kernels import AdditiveKernel, Kernel, MaternKernel, ScaleKernel
@@ -561,3 +563,19 @@ def posterior(
             **kwargs,
         )
         return GaussianMixturePosterior(distribution=posterior.distribution)
+
+    @classmethod
+    def construct_inputs(
+        cls,
+        training_data: SupervisedDataset,
+        *,
+        num_taus: int = 4,
+    ) -> dict[str, BotorchContainer | Tensor]:
+        r"""Construct `Model` keyword arguments from a dict of `SupervisedDataset`.
+
+        Args:
+            training_data: A `SupervisedDataset` containing the training data.
+            num_taus: Number of taus to use in the ensemble (4 if omitted).
+        """
+        base_inputs = super().construct_inputs(training_data=training_data)
+        return {**base_inputs, "num_taus": num_taus}

test/models/test_map_saas.py

@@ -34,6 +34,7 @@
 from botorch.posteriors.gpytorch import GPyTorchPosterior
 from botorch.test_utils.mock import mock_optimize
 from botorch.utils.constraints import LogTransformedInterval
+from botorch.utils.datasets import SupervisedDataset
 from botorch.utils.testing import BotorchTestCase
 from gpytorch.constraints import Interval
 from gpytorch.kernels import AdditiveKernel, MaternKernel, ScaleKernel
@@ -568,6 +569,56 @@ def test_ensemble_map_saas_validation(self) -> None:
                 train_X=torch.rand(2, 5, 3), train_Y=torch.rand(2, 5, 1)
             )
 
+    def test_ensemble_map_saas_construct_inputs(self) -> None:
+        """Test the construct_inputs class method for EnsembleMapSaasSingleTaskGP."""
+
+        train_X, train_Y, _ = self._get_data()
+        training_data = SupervisedDataset(
+            X=train_X, Y=train_Y, feature_names=["x1", "x2", "x3"], outcome_names=["y"]
+        )
+
+        # Test with default num_taus
+        inputs_default = EnsembleMapSaasSingleTaskGP.construct_inputs(
+            training_data=training_data
+        )
+        self.assertIn("num_taus", inputs_default)
+        self.assertEqual(inputs_default["num_taus"], 4)
+        self.assertIn("train_X", inputs_default)
+        self.assertIn("train_Y", inputs_default)
+        self.assertAllClose(inputs_default["train_X"], train_X)
+        self.assertAllClose(inputs_default["train_Y"], train_Y)
+
+        # Test with custom num_taus
+        custom_num_taus = 6
+        inputs_custom = EnsembleMapSaasSingleTaskGP.construct_inputs(
+            training_data=training_data, num_taus=custom_num_taus
+        )
+        self.assertIn("num_taus", inputs_custom)
+        self.assertEqual(inputs_custom["num_taus"], custom_num_taus)
+        self.assertIn("train_X", inputs_custom)
+        self.assertIn("train_Y", inputs_custom)
+        self.assertAllClose(inputs_custom["train_X"], train_X)
+        self.assertAllClose(inputs_custom["train_Y"], train_Y)
+
+        # Test with train_Yvar in the dataset
+        train_Yvar = 0.1 * torch.rand_like(train_Y)
+        training_data_with_yvar = SupervisedDataset(
+            X=train_X,
+            Y=train_Y,
+            Yvar=train_Yvar,
+            feature_names=["x1", "x2", "x3"],
+            outcome_names=["y"],
+        )
+        inputs_with_yvar = EnsembleMapSaasSingleTaskGP.construct_inputs(
+            training_data=training_data_with_yvar, num_taus=3
+        )
+        self.assertIn("train_Yvar", inputs_with_yvar)
+        self.assertAllClose(inputs_with_yvar["train_Yvar"], train_Yvar)
+        self.assertEqual(inputs_with_yvar["num_taus"], 3)
+        model_with_yvar = EnsembleMapSaasSingleTaskGP(**inputs_with_yvar)
+        self.assertIsInstance(model_with_yvar, EnsembleMapSaasSingleTaskGP)
+        self.assertEqual(model_with_yvar.batch_shape, torch.Size([3]))
+
 
 class TestAdditiveMapSaasSingleTaskGP(BotorchTestCase):
     def _get_data_and_model(