Support custom mean & likelihood in MultiTaskGP (pytorch#1909)

Summary:
Pull Request resolved: pytorch#1909

Title.

Reviewed By: esantorella

Differential Revision: D47130193

fbshipit-source-id: 16c78afcec7d8f27ff001317dbfedeb573dab61e

botorch/models/multitask.py

@@ -57,6 +57,7 @@
     FixedNoiseGaussianLikelihood,
     GaussianLikelihood,
 )
+from gpytorch.likelihoods.likelihood import Likelihood
 from gpytorch.likelihoods.multitask_gaussian_likelihood import (
     MultitaskGaussianLikelihood,
 )
@@ -106,7 +107,9 @@ def __init__(
         train_Y: Tensor,
         task_feature: int,
         train_Yvar: Optional[Tensor] = None,
+        mean_module: Optional[Module] = None,
         covar_module: Optional[Module] = None,
+        likelihood: Optional[Likelihood] = None,
         task_covar_prior: Optional[Prior] = None,
         output_tasks: Optional[List[int]] = None,
         rank: Optional[int] = None,
@@ -121,10 +124,16 @@ def __init__(
                 features (see `task_feature` argument).
             train_Y: A `n x 1` or `b x n x 1` (batch mode) tensor of training
                 observations.
+            task_feature: The index of the task feature (`-d <= task_feature <= d`).
             train_Yvar: An optional `n` or `b x n` (batch mode) tensor of observed
                 measurement noise. If None, we infer the noise.
                 Note that the inferred noise is common across all tasks.
-            task_feature: The index of the task feature (`-d <= task_feature <= d`).
+            mean_module: The mean function to be used. Defaults to `ConstantMean`.
+            covar_module: The module for computing the covariance matrix between
+                the non-task features. Defaults to `MaternKernel`.
+            likelihood: A likelihood. The default is selected based on `train_Yvar`.
+                If `train_Yvar` is None, a standard `GaussianLikelihood` with inferred
+                noise level is used. Otherwise, a FixedNoiseGaussianLikelihood is used.
             output_tasks: A list of task indices for which to compute model
                 outputs for. If omitted, return outputs for all task indices.
             rank: The rank to be used for the index kernel. If omitted, use a
@@ -170,10 +179,11 @@ def __init__(
         self._num_outputs = len(output_tasks)
 
         # TODO (T41270962): Support task-specific noise levels in likelihood
-        if train_Yvar is None:
-            likelihood = GaussianLikelihood(noise_prior=GammaPrior(1.1, 0.05))
-        else:
-            likelihood = FixedNoiseGaussianLikelihood(noise=train_Yvar.squeeze(-1))
+        if likelihood is None:
+            if train_Yvar is None:
+                likelihood = GaussianLikelihood(noise_prior=GammaPrior(1.1, 0.05))
+            else:
+                likelihood = FixedNoiseGaussianLikelihood(noise=train_Yvar.squeeze(-1))
 
         # construct indexer to be used in forward
         self._task_feature = task_feature
@@ -183,7 +193,7 @@ def __init__(
         super().__init__(
             train_inputs=train_X, train_targets=train_Y, likelihood=likelihood
         )
-        self.mean_module = ConstantMean()
+        self.mean_module = mean_module or ConstantMean()
         if covar_module is None:
             self.covar_module = get_matern_kernel_with_gamma_prior(
                 ard_num_dims=self.num_non_task_features

test/models/test_multitask.py

@@ -38,6 +38,7 @@
     MultitaskGaussianLikelihood,
 )
 from gpytorch.means import ConstantMean, MultitaskMean
+from gpytorch.means.linear_mean import LinearMean
 from gpytorch.mlls.exact_marginal_log_likelihood import ExactMarginalLogLikelihood
 from gpytorch.priors import GammaPrior, LogNormalPrior, SmoothedBoxPrior
 from gpytorch.priors.lkj_prior import LKJCovariancePrior
@@ -377,6 +378,21 @@ def test_MultiTaskGP_given_covar_module(self):
             self.assertAlmostEqual(model.covar_module.lengthscale_prior.loc, 0.0)
             self.assertAlmostEqual(model.covar_module.lengthscale_prior.scale, 1.0)
 
+    def test_custom_mean_and_likelihood(self):
+        tkwargs = {"device": self.device, "dtype": torch.double}
+        _, (train_X, train_Y) = _gen_datasets(**tkwargs)
+        mean_module = LinearMean(input_size=train_X.shape[-1])
+        likelihood = GaussianLikelihood(noise_prior=LogNormalPrior(0, 1))
+        model = MultiTaskGP(
+            train_X,
+            train_Y,
+            task_feature=0,
+            mean_module=mean_module,
+            likelihood=likelihood,
+        )
+        self.assertIs(model.mean_module, mean_module)
+        self.assertIs(model.likelihood, likelihood)
+
 
 class TestFixedNoiseMultiTaskGP(BotorchTestCase):
     def test_deprecation_warning(self):