From 8f88a89e0ce08db3fbacc341e85fbb55a442cb11 Mon Sep 17 00:00:00 2001
From: Max Balandat <balandat@meta.com>
Date: Sun, 15 Sep 2024 19:07:28 -0700
Subject: [PATCH] Account for botorch removal of deprecated models (#2769)

Summary:
Pull Request resolved: https://github.com/facebook/Ax/pull/2769

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

fbshipit-source-id: b51b2986fba07d10cda02941bcf95c64864b7da6
---
 ax/models/torch/botorch_defaults.py    | 2 +-
 ax/storage/botorch_modular_registry.py | 9 ++-------
 tutorials/sebo.ipynb                   | 4 ++--
 3 files changed, 5 insertions(+), 10 deletions(-)

diff --git a/ax/models/torch/botorch_defaults.py b/ax/models/torch/botorch_defaults.py
index 81397d4cf04..05e16b203a2 100644
--- 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
index 7dcc00c2cbc..178cb470e82 100644
--- 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
@@ -110,8 +107,6 @@
     # NOTE: Fixed noise models are deprecated. They point to their
     # 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
index 6c527f9e3a8..059b165b504 100644
--- 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",