Add construct_inputs to SAASBO (#1136)

Summary:
Pull Request resolved: #1136

Needed for BotAx support.

Reviewed By: saitcakmak

Differential Revision: D35164975

fbshipit-source-id: 2f496508c3f37a00c49fcbb971cf868740c15a22

Author: dme65

botorch/models/fully_bayesian.py

@@ -32,6 +32,7 @@
 from botorch.models.utils import validate_input_scaling
 from botorch.posteriors.fully_bayesian import FullyBayesianPosterior
 from botorch.sampling.samplers import MCSampler
+from botorch.utils.containers import TrainingData
 from gpytorch.constraints import GreaterThan
 from gpytorch.distributions.multivariate_normal import MultivariateNormal
 from gpytorch.kernels import MaternKernel, ScaleKernel
@@ -419,3 +420,19 @@ def posterior(
             marginalize_over_mcmc_samples=marginalize_over_mcmc_samples,
         )
         return posterior
+
+    @classmethod
+    def construct_inputs(
+        cls, training_data: TrainingData, **kwargs: Any
+    ) -> Dict[str, Any]:
+        r"""Construct kwargs for the `Model` from `TrainingData` and other options.
+
+        Args:
+            training_data: `TrainingData` container with data for single outcome
+                or for multiple outcomes for batched multi-output case.
+            **kwargs: None expected for this class.
+        """
+        inputs = {"train_X": training_data.X, "train_Y": training_data.Y}
+        if training_data.Yvar is not None:
+            inputs["train_Yvar"] = training_data.Yvar
+        return inputs

test/models/test_fully_bayesian.py

@@ -35,6 +35,7 @@
 from botorch.models.transforms import Normalize, Standardize
 from botorch.posteriors import FullyBayesianPosterior
 from botorch.sampling.samplers import IIDNormalSampler
+from botorch.utils.containers import TrainingData
 from botorch.utils.multi_objective.box_decompositions.non_dominated import (
     NondominatedPartitioning,
 )
@@ -360,3 +361,24 @@ def test_load_samples(self):
                         train_Yvar.clamp(MIN_INFERRED_NOISE_LEVEL),
                     )
                 )
+
+    def test_construct_inputs(self):
+        for infer_noise, dtype in itertools.product(
+            (True, False), (torch.float, torch.double)
+        ):
+            tkwargs = {"device": self.device, "dtype": dtype}
+            train_X, train_Y, train_Yvar, model = self._get_data_and_model(
+                infer_noise=infer_noise, **tkwargs
+            )
+            training_data = TrainingData.from_block_design(
+                X=train_X,
+                Y=train_Y,
+                Yvar=train_Yvar,
+            )
+            data_dict = model.construct_inputs(training_data)
+            if infer_noise:
+                self.assertTrue("train_Yvar" not in data_dict)
+            else:
+                self.assertTrue(torch.equal(data_dict["train_Yvar"], train_Yvar))
+            self.assertTrue(torch.equal(data_dict["train_X"], train_X))
+            self.assertTrue(torch.equal(data_dict["train_Y"], train_Y))