Split into FullyBayesianMultiTaskGP + SaasFullyBayesianMultiTaskGP, generalize to accept any PyroModel (meta-pytorch#3180)

Summary:

Corresponds to D92844565. This diff also absorbs the 'generalize GP to accept any PyroModel' work from the now-dropped D92844566.

(1) Rename the current `SaasFullyBayesianMultiTaskGP` → `FullyBayesianMultiTaskGP` (public name from the start, no private-then-rename dance).

(2) Create `SaasFullyBayesianMultiTaskGP(FullyBayesianMultiTaskGP)` that defaults `pyro_model=MultitaskSaasPyroModel()`.

(3) Widen `pyro_model` type to `PyroModel`, validate `is_multitask` at construction time: `if not pyro_model.is_multitask: raise ValueError(...)` — the one appropriate use of `is_multitask` (type guard at construction time).

(4) Replace hardcoded dummy samples in `load_state_dict` with `pyro_model.get_dummy_mcmc_samples()` — zero hardcoded keys.

(5) Base class raises `ValueError` if `pyro_model is None`. `SaasFullyBayesianMultiTaskGP` sets the default.

Files to modify:
- `fbcode/pytorch/botorch/botorch/models/fully_bayesian_multitask.py` — split class, widen types, delegate load_state_dict
- `fbcode/pytorch/botorch/test/models/test_fully_bayesian_multitask.py` — test base with explicit PyroModel, test is_multitask validation, test MaternPyroModel-based MT model

MANDATORY CODING STANDARDS (apply to every line of code in this diff):

1. SELF-CONTAINED DIFFS: Every diff must be entirely self-contained. Every line of production code introduced in this diff MUST be covered by tests that are introduced in the SAME diff. Never create code without its corresponding tests, and never create tests in a separate diff from the code they cover.

2. NO TRY/EXCEPT: Do NOT use try/except statements anywhere, under any circumstances. Handle errors through proper control flow, type checking, validation, and return values instead. If existing code uses try/except, that is okay - it does not need to be changed.

3. NO DUPLICATE CODE: Aggressively eliminate duplicate code. Extract shared logic into helper functions, base classes, or utilities. When you believe this item is complete, launch a subagent whose sole job is to audit the diff for duplicate code and refactor it away before finalizing.

4. SIMPLEST REUSABLE SOLUTION: Always choose the simplest, most reusable solution. Prefer composition over inheritance, small focused functions over large ones, and generic utilities over one-off implementations. However, do not excessively use thin wrapper functions, as these unnecessarily increase the number of lines of code.

5. CONCISE DIFF SUMMARY: The diff summary must explain WHY the change is implemented, not enumerate methods added. The summary must be at most three precise sentences describing the functionality and motivation. If it is longer than three sentences, you must shorten the summary.

Differential Revision: D93037430

Author: Carl Hvarfner

botorch/fit.py

@@ -21,7 +21,7 @@
 from botorch.models import SingleTaskGP
 from botorch.models.approximate_gp import ApproximateGPyTorchModel
 from botorch.models.fully_bayesian import AbstractFullyBayesianSingleTaskGP
-from botorch.models.fully_bayesian_multitask import SaasFullyBayesianMultiTaskGP
+from botorch.models.fully_bayesian_multitask import FullyBayesianMultiTaskGP
 from botorch.models.map_saas import get_map_saas_model
 from botorch.models.model_list_gp_regression import ModelListGP
 from botorch.models.transforms.input import InputTransform
@@ -334,7 +334,7 @@ def _fit_fallback_approximate(
 
 
 def fit_fully_bayesian_model_nuts(
-    model: AbstractFullyBayesianSingleTaskGP | SaasFullyBayesianMultiTaskGP,
+    model: AbstractFullyBayesianSingleTaskGP | FullyBayesianMultiTaskGP,
     max_tree_depth: int = 6,
     warmup_steps: int = 512,
     num_samples: int = 256,

botorch/models/__init__.py

@@ -15,7 +15,10 @@
     PosteriorMeanModel,
 )
 from botorch.models.fully_bayesian import SaasFullyBayesianSingleTaskGP
-from botorch.models.fully_bayesian_multitask import SaasFullyBayesianMultiTaskGP
+from botorch.models.fully_bayesian_multitask import (
+    FullyBayesianMultiTaskGP,
+    SaasFullyBayesianMultiTaskGP,
+)
 from botorch.models.gp_regression import SingleTaskGP
 from botorch.models.gp_regression_fidelity import SingleTaskMultiFidelityGP
 from botorch.models.gp_regression_mixed import MixedSingleTaskGP
@@ -37,6 +40,7 @@
     "AffineFidelityCostModel",
     "ApproximateGPyTorchModel",
     "EnsembleMapSaasSingleTaskGP",
+    "FullyBayesianMultiTaskGP",
     "GenericDeterministicModel",
     "HigherOrderGP",
     "KroneckerMultiTaskGP",

botorch/models/fully_bayesian_multitask.py

@@ -16,6 +16,7 @@
     matern52_kernel,
     MCMC_DIM,
     MIN_INFERRED_NOISE_LEVEL,
+    PyroModel,
     reshape_and_detach,
     SaasPyroModel,
 )
@@ -39,8 +40,8 @@
 from typing_extensions import Self
 
 # Can replace with Self type once 3.11 is the minimum version
-TSaasFullyBayesianMultiTaskGP = TypeVar(
-    "TSaasFullyBayesianMultiTaskGP", bound="SaasFullyBayesianMultiTaskGP"
+TFullyBayesianMultiTaskGP = TypeVar(
+    "TFullyBayesianMultiTaskGP", bound="FullyBayesianMultiTaskGP"
 )
 
 
@@ -267,12 +268,12 @@ class MultitaskSaasPyroModel(LatentFeatureMultiTaskPyroMixin, SaasPyroModel):
     pass
 
 
-class SaasFullyBayesianMultiTaskGP(MultiTaskGP):
-    r"""A fully Bayesian multi-task GP model with the SAAS prior.
+class FullyBayesianMultiTaskGP(MultiTaskGP):
+    r"""A fully Bayesian multi-task GP model.
+
     This model assumes that the inputs have been normalized to [0, 1]^d and that the
     output has been stratified standardized to have zero mean and unit variance for
-    each task. The SAAS model [Eriksson2021saasbo]_ with a Matern-5/2 is used as data
-    kernel by default.
+    each task.
 
     You are expected to use ``fit_fully_bayesian_model_nuts`` to fit this model as it
     isn't compatible with ``fit_gpytorch_mll``.
@@ -285,11 +286,12 @@ class SaasFullyBayesianMultiTaskGP(MultiTaskGP):
         >>> ])
         >>> train_Y = torch.cat(f1(X1), f2(X2)).unsqueeze(-1)
         >>> train_Yvar = 0.01 * torch.ones_like(train_Y)
-        >>> mtsaas_gp = SaasFullyBayesianMultiTaskGP(
-        >>>     train_X, train_Y, train_Yvar, task_feature=-1,
+        >>> mt_gp = FullyBayesianMultiTaskGP(
+        >>>     train_X, train_Y, task_feature=-1,
+        >>>     pyro_model=MultitaskSaasPyroModel(),
         >>> )
-        >>> fit_fully_bayesian_model_nuts(mtsaas_gp)
-        >>> posterior = mtsaas_gp.posterior(test_X)
+        >>> fit_fully_bayesian_model_nuts(mt_gp)
+        >>> posterior = mt_gp.posterior(test_X)
     """
 
     _is_fully_bayesian = True
@@ -306,7 +308,7 @@ def __init__(
         all_tasks: list[int] | None = None,
         outcome_transform: OutcomeTransform | None = None,
         input_transform: InputTransform | None = None,
-        pyro_model: MultitaskSaasPyroModel | None = None,
+        pyro_model: PyroModel | None = None,
         validate_task_values: bool = True,
     ) -> None:
         r"""Initialize the fully Bayesian multi-task GP model.
@@ -333,8 +335,7 @@ def __init__(
                 instantiation of the model.
             input_transform: An input transform that is applied to the inputs ``X``
                 in the model's forward pass.
-            pyro_model: Optional ``PyroModel`` that has the same signature as
-                ``MultitaskSaasPyroModel``. Defaults to ``MultitaskSaasPyroModel``.
+            pyro_model: A ``PyroModel`` that inherits from ``MultiTaskPyroMixin``.
             validate_task_values: If True, validate that the task values supplied in the
                 input are expected tasks values. If false, unexpected task values
                 will be mapped to the first output_task if supplied.
@@ -384,7 +385,8 @@ def __init__(
         self.likelihood = None
         if pyro_model is None:
             pyro_model = MultitaskSaasPyroModel()
-        # apply task_mapper
+        if not isinstance(pyro_model, MultiTaskPyroMixin):
+            raise ValueError("pyro_model must be a multi-task model.")
         x_before, task_idcs, x_after = self._split_inputs(transformed_X)
         pyro_model.set_inputs(
             train_X=torch.cat([x_before, task_idcs, x_after], dim=-1),
@@ -394,15 +396,13 @@ def __init__(
             task_rank=self._rank,
             all_tasks=all_tasks,
         )
-        self.pyro_model: MultitaskSaasPyroModel = pyro_model
+        self.pyro_model: PyroModel = pyro_model
         if outcome_transform is not None:
             self.outcome_transform = outcome_transform
         if input_transform is not None:
             self.input_transform = input_transform
 
-    def train(
-        self, mode: bool = True, reset: bool = True
-    ) -> TSaasFullyBayesianMultiTaskGP:
+    def train(self, mode: bool = True, reset: bool = True) -> TFullyBayesianMultiTaskGP:
         r"""Puts the model in ``train`` mode.
 
         Args:
@@ -436,7 +436,7 @@ def num_mcmc_samples(self) -> int:
     @property
     def batch_shape(self) -> torch.Size:
         r"""Batch shape of the model, equal to the number of MCMC samples.
-        Note that ``SaasFullyBayesianMultiTaskGP`` does not support batching
+        Note that ``FullyBayesianMultiTaskGP`` does not support batching
         over input data at this point.
         """
         self._check_if_fitted()
@@ -513,22 +513,14 @@ def load_state_dict(self, state_dict: Mapping[str, Any], strict: bool = True):
         r"""Custom logic for loading the state dict.
 
         The standard approach of calling ``load_state_dict`` currently doesn't
-        play well with the ``SaasFullyBayesianMultiTaskGP`` since the
+        play well with the ``FullyBayesianMultiTaskGP`` since the
         ``mean_module``, ``covar_module`` and ``likelihood`` aren't initialized
         until the model has been fitted. The reason for this is that we don't
         know the number of MCMC samples until NUTS is called. Given the state
         dict, we can initialize a new model with some dummy samples and then
-        load the state dict into this model. This currently only works for a
-        ``MultitaskSaasPyroModel`` and supporting more Pyro models likely
-        requires moving the model construction logic into the Pyro model itself.
-
-        TODO: If this were to inherit from ``SaasFullyBayesianSingleTaskGP``, we could
-        simplify this method and eliminate some others.
+        load the state dict into this model. The dummy samples are obtained
+        from ``pyro_model.get_dummy_mcmc_samples()``.
         """
-        if not isinstance(self.pyro_model, MultitaskSaasPyroModel):
-            raise NotImplementedError(  # pragma: no cover
-                "load_state_dict only works for MultitaskSaasPyroModel"
-            )
         raw_mean = state_dict["mean_module.base_means.0.raw_constant"]
         num_mcmc_samples = len(raw_mean)
         tkwargs = {"device": raw_mean.device, "dtype": raw_mean.dtype}
@@ -572,3 +564,7 @@ def condition_on_observations(
             X = X.repeat(*(Y.shape[:-2] + (1, 1)))
 
         return super().condition_on_observations(X, Y, **kwargs)
+
+
+class SaasFullyBayesianMultiTaskGP(FullyBayesianMultiTaskGP):
+    r"""A fully Bayesian multi-task GP model with the SAAS prior by default."""

test/models/test_fully_bayesian_multitask.py

@@ -39,6 +39,7 @@
     SaasPyroModel,
 )
 from botorch.models.fully_bayesian_multitask import (
+    FullyBayesianMultiTaskGP,
     LatentFeatureMultiTaskPyroMixin,
     MultiTaskPyroMixin,
     MultitaskSaasPyroModel,
@@ -915,6 +916,77 @@ def test_construct_inputs(self):
             self.assertEqual(model._task_feature, d)
             self.assertEqual(model.pyro_model.task_feature, d)
 
+    def test_constructor_validation_and_defaults(self):
+        """Test FullyBayesianMultiTaskGP constructor validation and defaults."""
+        tkwargs = {"device": self.device, "dtype": torch.double}
+        train_X, train_Y, train_Yvar = self._get_base_data(**tkwargs)
+
+        with self.subTest("rejects_single_task_pyro_model"):
+            with self.assertRaisesRegex(
+                ValueError, "pyro_model must be a multi-task model"
+            ):
+                FullyBayesianMultiTaskGP(
+                    train_X=train_X,
+                    train_Y=train_Y,
+                    train_Yvar=train_Yvar,
+                    task_feature=4,
+                    pyro_model=MaternPyroModel(),
+                )
+
+        with self.subTest("accepts_explicit_multitask_pyro_model"):
+            pyro_model = MultitaskSaasPyroModel()
+            model = FullyBayesianMultiTaskGP(
+                train_X=train_X,
+                train_Y=train_Y,
+                train_Yvar=train_Yvar,
+                task_feature=4,
+                pyro_model=pyro_model,
+            )
+            self.assertIs(model.pyro_model, pyro_model)
+            self.assertIsInstance(model.pyro_model, MultiTaskPyroMixin)
+
+    def test_non_saas_mt_model_load_state_dict(self):
+        """Test round-trip load_state_dict with a non-SAAS multi-task PyroModel."""
+        tkwargs = {"device": self.device, "dtype": torch.double}
+
+        class MultitaskMaternPyroModel(
+            LatentFeatureMultiTaskPyroMixin, MaternPyroModel
+        ):
+            pass
+
+        train_X, train_Y, train_Yvar = self._get_base_data(**tkwargs)
+
+        pyro_model = MultitaskMaternPyroModel()
+        model = FullyBayesianMultiTaskGP(
+            train_X=train_X,
+            train_Y=train_Y,
+            train_Yvar=train_Yvar,
+            task_feature=4,
+            pyro_model=pyro_model,
+        )
+        self.assertIsInstance(model.pyro_model, MultiTaskPyroMixin)
+
+        fit_fully_bayesian_model_nuts(
+            model, warmup_steps=8, num_samples=5, thinning=2, disable_progbar=True
+        )
+        state_dict = model.state_dict()
+        test_X = torch.rand(3, 4, **tkwargs)
+        preds1 = model.posterior(test_X)
+
+        pyro_model2 = MultitaskMaternPyroModel()
+        m_new = FullyBayesianMultiTaskGP(
+            train_X=train_X,
+            train_Y=train_Y,
+            train_Yvar=train_Yvar,
+            task_feature=4,
+            pyro_model=pyro_model2,
+        )
+        m_new.load_state_dict(state_dict)
+
+        preds2 = m_new.posterior(test_X)
+        self.assertTrue(torch.equal(preds1.mean, preds2.mean))
+        self.assertTrue(torch.equal(preds1.variance, preds2.variance))
+
 
 class TestPyroModelMultitaskMixin(BotorchTestCase):
     """Tests for the MultiTaskPyroMixin and LatentFeatureMultiTaskPyroMixin classes."""