Remove FixedNoiseMultiTaskGP

Summary: Deprecated since #1818. `MultiTaskGP` offers the full functionality of `FixedNoiseMultiTaskGP` without requiring a different class.

Reviewed By: dme65

Differential Revision: D56802249

fbshipit-source-id: 44b07d2a1b5551346871a9f8b36ca069b9280019

Author: saitcakmak

botorch/models/__init__.py

@@ -27,19 +27,14 @@
 from botorch.models.higher_order_gp import HigherOrderGP
 from botorch.models.model import ModelList
 from botorch.models.model_list_gp_regression import ModelListGP
-from botorch.models.multitask import (
-    FixedNoiseMultiTaskGP,
-    KroneckerMultiTaskGP,
-    MultiTaskGP,
-)
+from botorch.models.multitask import KroneckerMultiTaskGP, MultiTaskGP
 from botorch.models.pairwise_gp import PairwiseGP, PairwiseLaplaceMarginalLogLikelihood
 
 __all__ = [
     "AffineDeterministicModel",
     "AffineFidelityCostModel",
     "ApproximateGPyTorchModel",
     "FixedNoiseGP",
-    "FixedNoiseMultiTaskGP",
     "SaasFullyBayesianSingleTaskGP",
     "SaasFullyBayesianMultiTaskGP",
     "GenericDeterministicModel",

botorch/models/multitask.py

@@ -30,7 +30,6 @@
 from __future__ import annotations
 
 import math
-import warnings
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import torch
@@ -375,7 +374,12 @@ def construct_inputs(
 
         # Call Model.construct_inputs to parse training data
         base_inputs = super().construct_inputs(training_data=training_data)
-        if isinstance(training_data, MultiTaskDataset):
+        if (
+            isinstance(training_data, MultiTaskDataset)
+            # If task features are included in the data, all tasks will have
+            # some observations and they may have different task features.
+            and training_data.task_feature_index is None
+        ):
             all_tasks = list(range(len(training_data.datasets)))
             base_inputs["all_tasks"] = all_tasks
         if task_covar_prior is not None:
@@ -387,81 +391,6 @@ def construct_inputs(
         return base_inputs
 
 
-class FixedNoiseMultiTaskGP(MultiTaskGP):
-    r"""Multi-Task GP model using an ICM kernel, with known observation noise.
-
-    DEPRECATED: Please use `MultiTaskGP` with `train_Yvar` instead.
-    Will be removed in a future release (~v0.10).
-    """
-
-    def __init__(
-        self,
-        train_X: Tensor,
-        train_Y: Tensor,
-        train_Yvar: Tensor,
-        task_feature: int,
-        covar_module: Optional[Module] = None,
-        task_covar_prior: Optional[Prior] = None,
-        output_tasks: Optional[List[int]] = None,
-        rank: Optional[int] = None,
-        input_transform: Optional[InputTransform] = None,
-        outcome_transform: Optional[OutcomeTransform] = None,
-    ) -> None:
-        r"""
-        Args:
-            train_X: A `n x (d + 1)` or `b x n x (d + 1)` (batch mode) tensor
-                of training data. One of the columns should contain the task
-                features (see `task_feature` argument).
-            train_Y: A `n x 1` or `b x n x 1` (batch mode) tensor of training
-                observations.
-            train_Yvar: A `n` or `b x n` (batch mode) tensor of observed measurement
-                noise.
-            task_feature: The index of the task feature (`-d <= task_feature <= d`).
-            task_covar_prior : A Prior on the task covariance matrix. Must operate
-                on p.s.d. matrices. A common prior for this is the `LKJ` prior.
-            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
-                full rank (i.e. number of tasks) kernel.
-            input_transform: An input transform that is applied in the model's
-                forward pass.
-            outcome_transform: An outcome transform that is applied to the
-                training data during instantiation and to the posterior during
-                inference (that is, the `Posterior` obtained by calling
-                `.posterior` on the model will be on the original scale).
-
-        Example:
-            >>> X1, X2 = torch.rand(10, 2), torch.rand(20, 2)
-            >>> i1, i2 = torch.zeros(10, 1), torch.ones(20, 1)
-            >>> train_X = torch.cat([
-            >>>     torch.cat([X1, i1], -1), torch.cat([X2, i2], -1),
-            >>> ], dim=0)
-            >>> train_Y = torch.cat(f1(X1), f2(X2))
-            >>> train_Yvar = 0.1 + 0.1 * torch.rand_like(train_Y)
-            >>> model = FixedNoiseMultiTaskGP(train_X, train_Y, train_Yvar, -1)
-        """
-        warnings.warn(
-            "`FixedNoiseMultiTaskGP` has been deprecated and will be removed in a "
-            "future release. Please use the `MultiTaskGP` model instead. "
-            "When `train_Yvar` is specified, `MultiTaskGP` behaves the same "
-            "as the `FixedNoiseMultiTaskGP`.",
-            DeprecationWarning,
-            stacklevel=2,
-        )
-        super().__init__(
-            train_X=train_X,
-            train_Y=train_Y,
-            train_Yvar=train_Yvar,
-            covar_module=covar_module,
-            task_feature=task_feature,
-            output_tasks=output_tasks,
-            rank=rank,
-            task_covar_prior=task_covar_prior,
-            input_transform=input_transform,
-            outcome_transform=outcome_transform,
-        )
-
-
 class KroneckerMultiTaskGP(ExactGP, GPyTorchModel, FantasizeMixin):
     """Multi-task GP with Kronecker structure, using an ICM kernel.
 

botorch/utils/test_helpers.py

@@ -16,6 +16,7 @@
 
 import torch
 from botorch.acquisition.objective import PosteriorTransform
+from botorch.exceptions.errors import UnsupportedError
 from botorch.models.gpytorch import GPyTorchModel
 from botorch.models.model import FantasizeMixin, Model
 from botorch.models.transforms.outcome import Standardize
@@ -66,9 +67,24 @@ def standardize_moments(
 
 
 def gen_multi_task_dataset(
-    yvar: Optional[float] = None, task_values: Optional[List[int]] = None, **tkwargs
+    yvar: Optional[float] = None,
+    task_values: Optional[List[int]] = None,
+    skip_task_features_in_datasets: bool = False,
+    **tkwargs,
 ) -> Tuple[MultiTaskDataset, Tuple[Tensor, Tensor, Optional[Tensor]]]:
-    """Constructs a multi-task dataset with two tasks, each with 10 data points."""
+    """Constructs a multi-task dataset with two tasks, each with 10 data points.
+
+    Args:
+        yvar: The noise level to use for `train_Yvar`. If None, uses `train_Yvar=None`.
+        task_values: The values of the task features. If None, uses [0, 1].
+        skip_task_features_in_datasets: If True, the task features are not included in
+            Xs of the datasets used to construct the datasets. This is useful for
+            testing `MultiTaskDataset`.
+    """
+    if task_values is not None and skip_task_features_in_datasets:
+        raise UnsupportedError(  # pragma: no cover
+            "`task_values` and `skip_task_features_in_datasets` can't be used together."
+        )
     X = torch.linspace(0, 0.95, 10, **tkwargs) + 0.05 * torch.rand(10, **tkwargs)
     X = X.unsqueeze(dim=-1)
     Y1 = torch.sin(X * (2 * math.pi)) + torch.randn_like(X) * 0.2
@@ -81,24 +97,27 @@ def gen_multi_task_dataset(
     Yvar1 = None if yvar is None else torch.full_like(Y1, yvar)
     Yvar2 = None if yvar is None else torch.full_like(Y2, yvar)
     train_Yvar = None if yvar is None else torch.cat([Yvar1, Yvar2])
+    feature_slice = slice(1, None) if skip_task_features_in_datasets else slice(None)
     datasets = [
         SupervisedDataset(
-            X=train_X[:10],
+            X=train_X[:10, feature_slice],
             Y=Y1,
             Yvar=Yvar1,
-            feature_names=["task", "X"],
+            feature_names=["task", "X"][feature_slice],
             outcome_names=["y"],
         ),
         SupervisedDataset(
-            X=train_X[10:],
+            X=train_X[10:, feature_slice],
             Y=Y2,
             Yvar=Yvar2,
-            feature_names=["task", "X"],
+            feature_names=["task", "X"][feature_slice],
             outcome_names=["y1"],
         ),
     ]
     dataset = MultiTaskDataset(
-        datasets=datasets, target_outcome_name="y", task_feature_index=0
+        datasets=datasets,
+        target_outcome_name="y",
+        task_feature_index=None if skip_task_features_in_datasets else 0,
     )
     return dataset, (train_X, train_Y, train_Yvar)
 

test/models/test_contextual_multioutput.py

@@ -123,9 +123,14 @@ def test_FixedNoiseLCEMGP(self):
             self.assertIsInstance(model(test_x), MultivariateNormal)
 
     def test_construct_inputs(self) -> None:
-        for with_embedding_inputs, yvar in ((True, None), (False, 0.01)):
+        for with_embedding_inputs, yvar, skip_task_features_in_datasets in zip(
+            (True, False), (None, 0.01), (True, False), strict=True
+        ):
             dataset, (train_x, train_y, train_yvar) = gen_multi_task_dataset(
-                yvar=yvar, dtype=torch.double, device=self.device
+                yvar=yvar,
+                skip_task_features_in_datasets=skip_task_features_in_datasets,
+                dtype=torch.double,
+                device=self.device,
             )
             model_inputs = LCEMGP.construct_inputs(
                 training_data=dataset,
@@ -139,11 +144,18 @@ def test_construct_inputs(self) -> None:
                 ),
             )
             # Check that the model inputs are valid.
-            LCEMGP(**model_inputs)
+            model = LCEMGP(**model_inputs)
             # Check that the model inputs are as expected.
-            self.assertAllClose(model_inputs.pop("train_X"), train_x)
+            self.assertEqual(model.all_tasks, [0, 1])
+            if skip_task_features_in_datasets:
+                # In this case, the task feature is appended at the end.
+                self.assertAllClose(model_inputs.pop("train_X"), train_x[..., [1, 0]])
+                # all_tasks is inferred from data when task features are omitted.
+                self.assertEqual(model_inputs.pop("all_tasks"), [0, 1])
+            else:
+                self.assertAllClose(model_inputs.pop("train_X"), train_x)
             self.assertAllClose(model_inputs.pop("train_Y"), train_y)
-            if yvar is not None:
+            if train_yvar is not None:
                 self.assertAllClose(model_inputs.pop("train_Yvar"), train_yvar)
             if with_embedding_inputs:
                 self.assertEqual(model_inputs.pop("embs_dim_list"), [2])
@@ -155,7 +167,6 @@ def test_construct_inputs(self) -> None:
                     model_inputs.pop("context_cat_feature"),
                     torch.tensor([[0.4], [0.5]]),
                 )
-            self.assertEqual(model_inputs.pop("all_tasks"), [0, 1])
             self.assertEqual(model_inputs.pop("task_feature"), 0)
             self.assertIsNone(model_inputs.pop("output_tasks"))
             # Check that there are no unexpected inputs.