Add multi-task mixin support to PyroModel hierarchy

Summary: Generalize PyroModel with dispatch methods and multi-task mixins so that any PyroModel subclass can be composed with multi-task capabilities..

Differential Revision: D92844567

Author: Carl Hvarfner

botorch/models/fully_bayesian.py

@@ -165,14 +165,23 @@ def _maybe_input_warp(self, X: Tensor, **tkwargs: Any) -> Tensor:
             return self.train_X
 
     def set_inputs(
-        self, train_X: Tensor, train_Y: Tensor, train_Yvar: Tensor | None = None
+        self,
+        train_X: Tensor,
+        train_Y: Tensor,
+        train_Yvar: Tensor | None = None,
+        task_feature: int | None = None,
+        task_rank: int | None = None,
     ) -> None:
         """Set the training data.
 
         Args:
             train_X: Training inputs (n x d)
             train_Y: Training targets (n x 1)
             train_Yvar: Observed noise variance (n x 1). Inferred if None.
+            task_feature: The index of the task feature column. Used by
+                multi-task mixins.
+            task_rank: The number of learned task embeddings. Used by
+                multi-task mixins.
         """
         self.train_X = train_X
         self.train_Y = train_Y
@@ -192,6 +201,19 @@ def postprocess_mcmc_samples(
         """Post-process the final MCMC samples."""
         pass  # pragma: no cover
 
+    @abstractmethod
+    def get_dummy_mcmc_samples(
+        self,
+        num_mcmc_samples: int,
+        **tkwargs: Any,
+    ) -> dict[str, Tensor]:
+        """Return dummy MCMC samples for state dict loading.
+
+        Each subclass returns a dict of ones with the keys and shapes that
+        ``load_mcmc_samples`` expects.
+        """
+        pass  # pragma: no cover
+
     @abstractmethod
     def load_mcmc_samples(
         self, mcmc_samples: dict[str, Tensor]
@@ -245,6 +267,60 @@ def sample_concentrations(self, **tkwargs: Any) -> tuple[Tensor, Tensor]:
 
         return c0, c1
 
+    def _common_dummy_samples(
+        self,
+        mcmc_samples: dict[str, Tensor],
+        num_mcmc_samples: int,
+        **tkwargs: Any,
+    ) -> dict[str, Tensor]:
+        """Add noise and warping entries to ``mcmc_samples`` in-place."""
+        dim = self.ard_num_dims
+        if self.train_Yvar is None:
+            mcmc_samples["noise"] = torch.ones(num_mcmc_samples, **tkwargs)
+        if self.use_input_warping:
+            mcmc_samples["c0"] = torch.ones(num_mcmc_samples, dim, **tkwargs)
+            mcmc_samples["c1"] = torch.ones(num_mcmc_samples, dim, **tkwargs)
+        return mcmc_samples
+
+    def _prepare_features(self, X: Tensor, **tkwargs: Any) -> Tensor:
+        """Select feature columns for kernel computation.
+
+        Overridden by multi-task mixins to strip the task column.
+        """
+        return X
+
+    def _maybe_multitask_transform(
+        self, K_noiseless: Tensor, mean: Tensor, **tkwargs: Any
+    ) -> tuple[Tensor, Tensor]:
+        r"""Apply multi-task covariance transform to the kernel and mean.
+
+        No-op for single-task models. Overridden by multi-task mixins.
+        """
+        return K_noiseless, mean
+
+    def _preprocess_mean_for_loading(
+        self, mcmc_samples: dict[str, Tensor]
+    ) -> dict[str, Tensor]:
+        """Return mcmc_samples with a scalar mean for building covariance modules.
+
+        No-op for single-task models. Overridden by multi-task mixins.
+        """
+        return mcmc_samples
+
+    def _finalize_multitask_modules(
+        self,
+        mcmc_samples: dict[str, Tensor],
+        mean_module: Mean,
+        covar_module: Kernel,
+        batch_shape: torch.Size,
+        **tkwargs: Any,
+    ) -> tuple[Mean, Kernel]:
+        """Optionally replace mean/covar with multi-task versions.
+
+        No-op for single-task models. Overridden by multi-task mixins.
+        """
+        return mean_module, covar_module
+
     def sample_observations(
         self,
         mean: Tensor,
@@ -322,7 +398,11 @@ def sample(self) -> None:
         noise = self.sample_noise(**tkwargs)
         lengthscale = self.sample_lengthscale(dim=self.ard_num_dims, **tkwargs)
         X_tf = self._maybe_input_warp(self.train_X, **tkwargs)
+        X_tf = self._prepare_features(X_tf, **tkwargs)
         K_noiseless = outputscale * matern52_kernel(X=X_tf, lengthscale=lengthscale)
+        K_noiseless, mean = self._maybe_multitask_transform(
+            K_noiseless, mean, **tkwargs
+        )
         self.sample_observations(
             mean=mean, K_noiseless=K_noiseless, noise=noise, **tkwargs
         )
@@ -375,6 +455,18 @@ def postprocess_mcmc_samples(
         """
         return mcmc_samples
 
+    def get_dummy_mcmc_samples(
+        self,
+        num_mcmc_samples: int,
+        **tkwargs: Any,
+    ) -> dict[str, Tensor]:
+        """Return dummy MCMC samples for state dict loading."""
+        mcmc_samples = {
+            "mean": torch.ones(num_mcmc_samples, **tkwargs),
+            "lengthscale": torch.ones(num_mcmc_samples, self.ard_num_dims, **tkwargs),
+        }
+        return self._common_dummy_samples(mcmc_samples, num_mcmc_samples, **tkwargs)
+
     def _get_covar_module(
         self,
         use_scale_kernel: bool,
@@ -410,8 +502,10 @@ def load_mcmc_samples(
         num_mcmc_samples = len(mcmc_samples["mean"])
         batch_shape = torch.Size([num_mcmc_samples])
 
+        mean_mcmc = self._preprocess_mean_for_loading(mcmc_samples)
+
         mean_module = ConstantMean(batch_shape=batch_shape).to(**tkwargs)
-        outputscale = mcmc_samples.get("outputscale")
+        outputscale = mean_mcmc.get("outputscale")
         covar_module = self._get_covar_module(
             use_scale_kernel=outputscale is not None, batch_shape=batch_shape, **tkwargs
         )
@@ -446,7 +540,7 @@ def load_mcmc_samples(
         )
         mean_module.constant.data = reshape_and_detach(
             target=mean_module.constant.data,
-            new_value=mcmc_samples["mean"],
+            new_value=mean_mcmc["mean"],
         )
         if self.use_input_warping:
             indices = (
@@ -470,6 +564,14 @@ def load_mcmc_samples(
             )
         else:
             warping_function = None
+
+        mean_module, covar_module = self._finalize_multitask_modules(
+            mcmc_samples=mcmc_samples,
+            mean_module=mean_module,
+            covar_module=covar_module,
+            batch_shape=batch_shape,
+            **tkwargs,
+        )
         return mean_module, covar_module, likelihood, warping_function
 
 
@@ -529,6 +631,18 @@ def postprocess_mcmc_samples(
         del mcmc_samples["kernel_tausq"], mcmc_samples["_kernel_inv_length_sq"]
         return mcmc_samples
 
+    def get_dummy_mcmc_samples(
+        self,
+        num_mcmc_samples: int,
+        **tkwargs: Any,
+    ) -> dict[str, Tensor]:
+        """Return dummy MCMC samples for state dict loading."""
+        mcmc_samples = super().get_dummy_mcmc_samples(
+            num_mcmc_samples=num_mcmc_samples, **tkwargs
+        )
+        mcmc_samples["outputscale"] = torch.ones(num_mcmc_samples, **tkwargs)
+        return mcmc_samples
+
 
 class LinearPyroModel(PyroModel):
     r"""Implementation of a Bayesian Linear pyro model.
@@ -546,9 +660,13 @@ def sample(self) -> None:
         mean = self.sample_mean(**tkwargs)
         weight_variance = self.sample_weight_variance(**tkwargs)
         X_tf = self._maybe_input_warp(X=self.train_X, **tkwargs)
+        X_tf = self._prepare_features(X_tf, **tkwargs)
         X_tf = X_tf - 0.5  # center transformed data at 0 (for linear model)
         K_noiseless = linear_kernel(X=X_tf, weight_variance=weight_variance)
         noise = self.sample_noise(**tkwargs)
+        K_noiseless, mean = self._maybe_multitask_transform(
+            K_noiseless, mean, **tkwargs
+        )
         self.sample_observations(
             mean=mean, K_noiseless=K_noiseless, noise=noise, **tkwargs
         )
@@ -589,6 +707,20 @@ def postprocess_mcmc_samples(
         del mcmc_samples["tau_sq"], mcmc_samples["_weight_variance_sq"]
         return mcmc_samples
 
+    def get_dummy_mcmc_samples(
+        self,
+        num_mcmc_samples: int,
+        **tkwargs: Any,
+    ) -> dict[str, Tensor]:
+        """Return dummy MCMC samples for state dict loading."""
+        mcmc_samples = {
+            "mean": torch.ones(num_mcmc_samples, **tkwargs),
+            "weight_variance": torch.ones(
+                num_mcmc_samples, self.ard_num_dims, **tkwargs
+            ),
+        }
+        return self._common_dummy_samples(mcmc_samples, num_mcmc_samples, **tkwargs)
+
     def load_mcmc_samples(
         self, mcmc_samples: dict[str, Tensor]
     ) -> tuple[Mean, Kernel, Likelihood, InputTransform]:
@@ -975,27 +1107,6 @@ def median_lengthscale(self) -> Tensor:
         lengthscale = base_kernel.lengthscale.clone()
         return lengthscale.median(0).values.squeeze(0)
 
-    def _get_dummy_mcmc_samples(
-        self,
-        num_mcmc_samples: int,
-        dim: int,
-        dtype: torch.dtype,
-        device: torch.device,
-    ) -> dict[str, Tensor]:
-        # Load some dummy samples
-        tkwargs = {"dtype": dtype, "device": device}
-        mcmc_samples = {
-            "mean": torch.ones(num_mcmc_samples, **tkwargs),
-            "lengthscale": torch.ones(num_mcmc_samples, dim, **tkwargs),
-        }
-        if self.pyro_model.train_Yvar is None:
-            mcmc_samples["noise"] = torch.ones(num_mcmc_samples, **tkwargs)
-
-        if self.pyro_model.use_input_warping:
-            mcmc_samples["c0"] = torch.ones(num_mcmc_samples, dim, **tkwargs)
-            mcmc_samples["c1"] = torch.ones(num_mcmc_samples, dim, **tkwargs)
-        return mcmc_samples
-
     def load_state_dict(
         self, state_dict: Mapping[str, Any], strict: bool = True
     ) -> None:
@@ -1012,11 +1123,10 @@ def load_state_dict(
         the model construction logic into the Pyro model itself.
         """
         raw_mean = state_dict["mean_module.raw_constant"]
-        mcmc_samples = self._get_dummy_mcmc_samples(
-            num_mcmc_samples=len(raw_mean),
-            dim=self.pyro_model.train_X.shape[-1],
-            dtype=raw_mean.dtype,
-            device=raw_mean.device,
+        num_mcmc_samples = len(raw_mean)
+        tkwargs = {"dtype": raw_mean.dtype, "device": raw_mean.device}
+        mcmc_samples = self.pyro_model.get_dummy_mcmc_samples(
+            num_mcmc_samples=num_mcmc_samples, **tkwargs
         )
         self.load_mcmc_samples(mcmc_samples=mcmc_samples)
         # Load the actual samples from the state dict
@@ -1043,22 +1153,6 @@ class SaasFullyBayesianSingleTaskGP(FullyBayesianSingleTaskGP):
 
     _pyro_model_class: type[PyroModel] = SaasPyroModel
 
-    def _get_dummy_mcmc_samples(
-        self,
-        num_mcmc_samples: int,
-        dim: int,
-        dtype: torch.dtype,
-        device: torch.device,
-    ) -> dict[str, Tensor]:
-        mcmc_samples = super()._get_dummy_mcmc_samples(
-            num_mcmc_samples=num_mcmc_samples, dim=dim, dtype=dtype, device=device
-        )
-        # add outputscale
-        mcmc_samples["outputscale"] = torch.ones(
-            num_mcmc_samples, dtype=dtype, device=device
-        )
-        return mcmc_samples
-
 
 class FullyBayesianLinearSingleTaskGP(AbstractFullyBayesianSingleTaskGP):
     r"""A fully Bayesian single-task GP model with a linear kernel.
@@ -1102,19 +1196,10 @@ def load_state_dict(
         """
         weight_variance = state_dict["covar_module.raw_variance"]
         num_mcmc_samples = len(weight_variance)
-        dim = self.pyro_model.train_X.shape[-1]
         tkwargs = {"device": weight_variance.device, "dtype": weight_variance.dtype}
-        # Load some dummy samples
-        # deal with c0 c1
-        mcmc_samples = {
-            "mean": torch.ones(num_mcmc_samples, **tkwargs),
-            "weight_variance": torch.ones(num_mcmc_samples, dim, **tkwargs),
-        }
-        if self.pyro_model.use_input_warping:
-            mcmc_samples["c0"] = torch.ones(num_mcmc_samples, dim, **tkwargs)
-            mcmc_samples["c1"] = torch.ones(num_mcmc_samples, dim, **tkwargs)
-        if self.pyro_model.train_Yvar is None:
-            mcmc_samples["noise"] = torch.ones(num_mcmc_samples, **tkwargs)
+        mcmc_samples = self.pyro_model.get_dummy_mcmc_samples(
+            num_mcmc_samples=num_mcmc_samples, **tkwargs
+        )
         self.load_mcmc_samples(mcmc_samples=mcmc_samples)
         # Load the actual samples from the state dict
         super().load_state_dict(state_dict=state_dict, strict=strict)