add construct inputs for SaasFullyBayesianMultiTaskGP (#1203)

Summary:
Pull Request resolved: #1203

as title.

D35573439 is too big. I make a separate diff to add construct_inputs for using `SaasFullyBayesianMultiTaskGP`. This overwritten is necessary in order to use the model in BotAx

Reviewed By: dme65

Differential Revision: D36136031

fbshipit-source-id: c6aee985f7f5468252909306cf4f577818debcb0

Author: qingfeng10

botorch/models/fully_bayesian_multitask.py

@@ -25,6 +25,7 @@
 from botorch.models.transforms.outcome import OutcomeTransform
 from botorch.posteriors.fully_bayesian import FullyBayesianPosterior, MCMC_DIM
 from botorch.sampling.samplers import MCSampler
+from botorch.utils.datasets import SupervisedDataset
 from gpytorch.distributions.multivariate_normal import MultivariateNormal
 from gpytorch.kernels import MaternKernel
 from gpytorch.kernels.kernel import Kernel
@@ -372,3 +373,26 @@ def forward(self, X: Tensor) -> MultivariateNormal:
         covar_i = self.task_covar_module(latent_features)
         covar = covar_x.mul(covar_i)
         return MultivariateNormal(mean_x, covar)
+
+    @classmethod
+    def construct_inputs(
+        cls,
+        training_data: Dict[str, SupervisedDataset],
+        task_feature: int,
+        rank: Optional[int] = None,
+        **kwargs: Any,
+    ) -> Dict[str, Any]:
+        r"""Construct `Model` keyword arguments from dictionary of `SupervisedDataset`.
+
+        Args:
+            training_data: Dictionary of `SupervisedDataset`.
+            task_feature: Column index of embedded task indicator features. For details,
+                see `parse_training_data`.
+            rank: The rank of the cross-task covariance matrix.
+        """
+
+        inputs = super().construct_inputs(
+            training_data=training_data, task_feature=task_feature, rank=rank, **kwargs
+        )
+        inputs.pop("task_covar_prior")
+        return inputs

test/models/test_fully_bayesian_multitask.py

@@ -46,6 +46,8 @@
 from gpytorch.likelihoods import FixedNoiseGaussianLikelihood
 from gpytorch.means import ConstantMean
 
+from .test_multitask import _gen_fixed_noise_model_and_data
+
 
 class TestFullyBayesianMultiTaskGP(BotorchTestCase):
     def _get_data_and_model(self, task_rank: Optional[int] = 1, **tkwargs):
@@ -509,3 +511,33 @@ def test_load_samples(self):
                     mcmc_samples["latent_features"],
                 )
             )
+
+    def test_construct_inputs(self):
+        for dtype in [torch.float, torch.double]:
+            tkwargs = {"device": self.device, "dtype": dtype}
+            task_feature = 0
+
+            (
+                _,
+                datasets,
+                (train_X, train_Y, train_Yvar),
+            ) = _gen_fixed_noise_model_and_data(task_feature=task_feature, **tkwargs)
+
+            model = SaasFullyBayesianMultiTaskGP(
+                train_X=train_X,
+                train_Y=train_Y,
+                train_Yvar=train_Yvar,
+                task_feature=task_feature,
+            )
+
+            data_dict = model.construct_inputs(
+                datasets,
+                task_feature=task_feature,
+                rank=1,
+            )
+            self.assertTrue(torch.equal(data_dict["train_X"], train_X))
+            self.assertTrue(torch.equal(data_dict["train_Y"], train_Y))
+            self.assertTrue(torch.allclose(data_dict["train_Yvar"], train_Yvar))
+            self.assertEqual(data_dict["task_feature"], task_feature)
+            self.assertEqual(data_dict["rank"], 1)
+            self.assertTrue("task_covar_prior" not in data_dict)