Implement multitask SAAS (#1181)

Summary:
Pull Request resolved: #1181

Move multitask SAAS GP into BoTorch

Reviewed By: dme65

Differential Revision: D35573439

fbshipit-source-id: 92bef5e423e48092f8848562cf6495729c403c8f

Author: qingfeng10

botorch/fit.py

@@ -13,12 +13,14 @@
 import logging
 import warnings
 from copy import deepcopy
-from typing import Any, Callable
+from typing import Any, Callable, Union
 
 from botorch.exceptions.errors import UnsupportedError
 from botorch.exceptions.warnings import BotorchWarning, OptimizationWarning
 from botorch.models.converter import batched_to_model_list, model_list_to_batched
 from botorch.models.fully_bayesian import SaasFullyBayesianSingleTaskGP
+from botorch.models.fully_bayesian_multitask import SaasFullyBayesianMultiTaskGP
+
 from botorch.models.gpytorch import BatchedMultiOutputGPyTorchModel
 from botorch.optim.fit import fit_gpytorch_scipy
 from botorch.optim.utils import sample_all_priors
@@ -157,7 +159,7 @@ def fit_gpytorch_model(
 
 
 def fit_fully_bayesian_model_nuts(
-    model: SaasFullyBayesianSingleTaskGP,
+    model: Union[SaasFullyBayesianSingleTaskGP, SaasFullyBayesianMultiTaskGP],
     max_tree_depth: int = 6,
     warmup_steps: int = 512,
     num_samples: int = 256,

botorch/models/__init__.py

@@ -15,6 +15,8 @@
     PosteriorMeanModel,
 )
 from botorch.models.fully_bayesian import SaasFullyBayesianSingleTaskGP
+from botorch.models.fully_bayesian_multitask import SaasFullyBayesianMultiTaskGP
+
 from botorch.models.gp_regression import (
     FixedNoiseGP,
     HeteroskedasticSingleTaskGP,
@@ -39,6 +41,7 @@
     "FixedNoiseGP",
     "FixedNoiseMultiTaskGP",
     "SaasFullyBayesianSingleTaskGP",
+    "SaasFullyBayesianMultiTaskGP",
     "GenericDeterministicModel",
     "HeteroskedasticSingleTaskGP",
     "HigherOrderGP",

botorch/models/fully_bayesian.py

@@ -145,6 +145,12 @@ class SaasPyroModel(PyroModel):
     `covar_module`).
     """
 
+    def set_inputs(
+        self, train_X: Tensor, train_Y: Tensor, train_Yvar: Optional[Tensor] = None
+    ):
+        super().set_inputs(train_X, train_Y, train_Yvar)
+        self.ard_num_dims = self.train_X.shape[-1]
+
     def sample(self) -> None:
         r"""Sample from the SAAS model.
 
@@ -155,7 +161,7 @@ def sample(self) -> None:
         outputscale = self.sample_outputscale(concentration=2.0, rate=0.15, **tkwargs)
         mean = self.sample_mean(**tkwargs)
         noise = self.sample_noise(**tkwargs)
-        lengthscale = self.sample_lengthscale(dim=self.train_X.shape[-1], **tkwargs)
+        lengthscale = self.sample_lengthscale(dim=self.ard_num_dims, **tkwargs)
         k = matern52_kernel(X=self.train_X, lengthscale=lengthscale)
         k = outputscale * k + noise * torch.eye(self.train_X.shape[0], **tkwargs)
         pyro.sample(
@@ -252,7 +258,7 @@ def load_mcmc_samples(
         mean_module = ConstantMean(batch_shape=batch_shape).to(**tkwargs)
         covar_module = ScaleKernel(
             base_kernel=MaternKernel(
-                ard_num_dims=self.train_X.shape[-1],
+                ard_num_dims=self.ard_num_dims,
                 batch_shape=batch_shape,
             ),
             batch_shape=batch_shape,