update dispatch/storage for new Fully Bayesian models/refactor (facebook#3806)

Summary:
X-link: pytorch/botorch#2857


see title

Reviewed By: Balandat

Differential Revision: D74824675

Author: sdaulton

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/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",
 }