Generic FullyBayesianMultiTaskGP

Summary: Creates generic FullyBayesianMultiTaskGP implementation and refactors the existing SAAS model as a special case

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,11 @@
     PosteriorMeanModel,
 )
 from botorch.models.fully_bayesian import SaasFullyBayesianSingleTaskGP
-from botorch.models.fully_bayesian_multitask import SaasFullyBayesianMultiTaskGP
+from botorch.models.fully_bayesian_multitask import (
+    FullyBayesianMultiTaskGP,
+    PositiveCorrelationSaasPyroModel,
+    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 +41,7 @@
     "AffineFidelityCostModel",
     "ApproximateGPyTorchModel",
     "EnsembleMapSaasSingleTaskGP",
+    "FullyBayesianMultiTaskGP",
     "GenericDeterministicModel",
     "HigherOrderGP",
     "KroneckerMultiTaskGP",
@@ -46,6 +51,7 @@
     "MultiTaskGP",
     "PairwiseGP",
     "PairwiseLaplaceMarginalLogLikelihood",
+    "PositiveCorrelationSaasPyroModel",
     "PosteriorMeanModel",
     "SaasFullyBayesianMultiTaskGP",
     "SaasFullyBayesianSingleTaskGP",