add option for input warping in non-linear fully Bayesian GPs

ax/generators/torch/botorch_modular/utils.py

@@ -32,7 +32,10 @@
 )
 from botorch.fit import fit_fully_bayesian_model_nuts, fit_gpytorch_mll
 from botorch.models import PairwiseLaplaceMarginalLogLikelihood
-from botorch.models.fully_bayesian import FullyBayesianSingleTaskGP
+from botorch.models.fully_bayesian import (
+    AbstractFullyBayesianSingleTaskGP,
+    FullyBayesianSingleTaskGP,
+)
 from botorch.models.fully_bayesian_multitask import SaasFullyBayesianMultiTaskGP
 from botorch.models.gp_regression import SingleTaskGP
 from botorch.models.gp_regression_fidelity import SingleTaskMultiFidelityGP
@@ -532,10 +535,11 @@ def _fit_botorch_model_gpytorch(
     fit_gpytorch_mll(mll)
 
 
-@fit_botorch_model.register(FullyBayesianSingleTaskGP)
-@fit_botorch_model.register(SaasFullyBayesianMultiTaskGP)
+@fit_botorch_model.register(
+    (AbstractFullyBayesianSingleTaskGP, SaasFullyBayesianMultiTaskGP)
+)
 def _fit_botorch_model_fully_bayesian_nuts(
-    model: FullyBayesianSingleTaskGP | SaasFullyBayesianMultiTaskGP,
+    model: AbstractFullyBayesianSingleTaskGP | SaasFullyBayesianMultiTaskGP,
     mll_class: type[MarginalLogLikelihood],
     mll_options: dict[str, Any] | None = None,
 ) -> None:

ax/generators/torch/tests/test_surrogate.py

@@ -656,7 +656,10 @@ def test_construct_model(self) -> None:
                 self.assertTrue(torch.equal(call_kwargs["train_Y"], self.Ys[0]))
                 self.assertIsInstance(call_kwargs["input_transform"], Normalize)
                 self.assertIsNone(call_kwargs["outcome_transform"])
-                self.assertEqual(len(call_kwargs), 4)
+                self.assertEqual(
+                    len(call_kwargs),
+                    6 if botorch_model_class is SaasFullyBayesianSingleTaskGP else 4,
+                )
 
                 mock_construct_inputs.assert_called_with(
                     training_data=self.training_data[0],

ax/storage/botorch_modular_registry.py

@@ -56,9 +56,12 @@
     AnalyticExpectedUtilityOfBestOption,
     qExpectedUtilityOfBestOption,
 )
-from botorch.models import SaasFullyBayesianSingleTaskGP
 from botorch.models.contextual import LCEAGP
-from botorch.models.fully_bayesian import FullyBayesianLinearSingleTaskGP
+from botorch.models.fully_bayesian import (
+    FullyBayesianLinearSingleTaskGP,
+    FullyBayesianSingleTaskGP,
+    SaasFullyBayesianSingleTaskGP,
+)
 from botorch.models.fully_bayesian_multitask import SaasFullyBayesianMultiTaskGP
 
 # BoTorch `Model` imports
@@ -124,6 +127,7 @@
     SingleTaskMultiFidelityGP: "SingleTaskMultiFidelityGP",
     FullyBayesianLinearSingleTaskGP: "FullyBayesianLinearSingleTaskGP",
     SaasFullyBayesianSingleTaskGP: "SaasFullyBayesianSingleTaskGP",
+    FullyBayesianSingleTaskGP: "FullyBayesianSingleTaskGP",
     SaasFullyBayesianMultiTaskGP: "SaasFullyBayesianMultiTaskGP",
     LCEAGP: "LCEAGP",
 }