Account for botorch removal of deprecated models

Summary: Account for removal of `FixedNoiseGP` and `FixedNoiseMultiFidelityGP` models in botorch, as well as the removal of the `data_fidelity` argument of `SingleTaskMultiFidelityGP`. Reviewed By: saitcakmak Differential Revision: D62687970
facebook · Sep 14, 2024 · ffb97dc · ffb97dc
1 parent 2fc80f1
commit ffb97dc
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 9 deletions.
diff --git a/ax/models/torch/botorch_defaults.py b/ax/models/torch/botorch_defaults.py
@@ -761,7 +761,7 @@ def _get_model(
         gp = SingleTaskMultiFidelityGP(
             train_X=X,
             train_Y=Y,
-            data_fidelity=fidelity_features[0],
+            data_fidelities=fidelity_features[:1],
             input_transform=warp_tf,
             **kwargs,
         )

diff --git a/ax/storage/botorch_modular_registry.py b/ax/storage/botorch_modular_registry.py
@@ -54,11 +54,8 @@
 from botorch.models.fully_bayesian_multitask import SaasFullyBayesianMultiTaskGP
 
 # BoTorch `Model` imports
-from botorch.models.gp_regression import FixedNoiseGP, SingleTaskGP
-from botorch.models.gp_regression_fidelity import (
-    FixedNoiseMultiFidelityGP,
-    SingleTaskMultiFidelityGP,
-)
+from botorch.models.gp_regression import SingleTaskGP
+from botorch.models.gp_regression_fidelity import SingleTaskMultiFidelityGP
 from botorch.models.gp_regression_mixed import MixedSingleTaskGP
 from botorch.models.model import Model
 from botorch.models.model_list_gp_regression import ModelListGP
@@ -111,7 +108,6 @@
     # supported parent classes, so that we can reap them with minimal
     # concern for backwards compatibility when the time comes.
     FixedNoiseGP: "SingleTaskGP",
-    FixedNoiseMultiFidelityGP: "SingleTaskMultiFidelityGP",
     MixedSingleTaskGP: "MixedSingleTaskGP",
     ModelListGP: "ModelListGP",
     MultiTaskGP: "MultiTaskGP",

diff --git a/tutorials/sebo.ipynb b/tutorials/sebo.ipynb
@@ -61,7 +61,7 @@
     "from ax.service.ax_client import AxClient, ObjectiveProperties\n",
     "from ax.utils.common.typeutils import checked_cast\n",
     "from botorch.acquisition.multi_objective import qNoisyExpectedHypervolumeImprovement\n",
-    "from botorch.models import FixedNoiseGP, SaasFullyBayesianSingleTaskGP, SingleTaskGP"
+    "from botorch.models import SaasFullyBayesianSingleTaskGP, SingleTaskGP"
    ]
   },
   {
@@ -247,7 +247,7 @@
     "if SMOKE_TEST:\n",
     "    N_BATCHES = 1\n",
     "    BATCH_SIZE = 1\n",
-    "    SURROGATE_CLASS = None  # Auto-pick SingleTaskGP / FixedNoiseGP\n",
+    "    SURROGATE_CLASS = None  # Auto-pick SingleTaskGP\n",
     "else:\n",
     "    N_BATCHES = 4\n",
     "    BATCH_SIZE = 5\n",