Back out "Enabled posterior samples as SurrogateBenchmarks"

mpolson64 · facebook-github-bot · commit eb9ab4d4d82e · 2025-09-05T11:36:04.000-07:00
Summary:
Original commit changeset: b5be2d69d998

Original Phabricator Diff: D80347570

Same motivation as D81695384, will be cleaned up after Ax 1.1.1 release

Differential Revision: D81799749
diff --git a/ax/benchmark/benchmark_test_functions/surrogate.py b/ax/benchmark/benchmark_test_functions/surrogate.py
@@ -15,17 +15,9 @@
 from ax.core.types import TParamValue
 from ax.utils.common.base import Base
 from ax.utils.common.equality import equality_typechecker
-from botorch.models.deterministic import (
-    DeterministicModel,
-    MatheronPathModel,
-    PosteriorMeanModel,
-)
-from botorch.utils.transforms import is_ensemble
 from pyre_extensions import none_throws
 from torch import Tensor
 
-RANDOM_SURROGATE_TYPES: list[type[DeterministicModel]] = [MatheronPathModel]
-
 
 @dataclass(kw_only=True)
 class SurrogateTestFunction(BenchmarkTestFunction):
@@ -43,21 +35,12 @@ class SurrogateTestFunction(BenchmarkTestFunction):
         get_surrogate: Function that returns the surrogate, to allow for lazy
             construction. If `get_surrogate` is not provided, `surrogate` must
             be provided and vice versa.
-        surrogate_model_type: The type of surrogate model to use. We either pass
-            in a type of deterministic model, (e.g. PosteriorMeanModel) or
-            a function that returns a deterministic model
-            (e.g. get_matheron_path_model).
-        sample_from_ensemble: If True, when the surrogate is an ensemble model,
-            we sample from the ensemble instead of averaging over it.
     """
 
     name: str
     outcome_names: Sequence[str]
     _surrogate: TorchAdapter | None = None
     get_surrogate: None | Callable[[], TorchAdapter] = None
-    surrogate_model_type: type[DeterministicModel] = PosteriorMeanModel
-    sample_from_ensemble: bool = False
-    seed: int = 0
 
     def __post_init__(self) -> None:
         if self.get_surrogate is None and self._surrogate is None:
@@ -66,50 +49,10 @@ def __post_init__(self) -> None:
                 " vice versa."
             )
 
-    def wrap_surrogate_in_deterministic_model(self) -> None:
-        """Substitute the surrogate model for a deterministic model used for
-        benchmarking."""
-        # pyre-ignore[16]: `ax.generators.torch_base.TorchGenerator` has no attribute
-        # `surrogate`.
-        surrogate_model = none_throws(self._surrogate).generator.surrogate
-
-        if isinstance(surrogate_model.model, DeterministicModel):
-            return  # Already wrapped
-
-        base_model = surrogate_model.model
-
-        # Check if surrogate_model_type accepts a 'seed' argument
-        if self.surrogate_model_type in RANDOM_SURROGATE_TYPES:
-            # pyre-ignore[28]: Unexpected keyword argument `seed` to call
-            # `botorch.models.ensemble.EnsembleModel.__init__`.
-            wrapped_model = self.surrogate_model_type(base_model, seed=self.seed)
-        else:
-            wrapped_model = self.surrogate_model_type(base_model)
-
-        # If the model is an ensemble, we want to choose one of the models instead
-        # of averging over them.
-        if is_ensemble(base_model) and self.sample_from_ensemble:
-            num_models = base_model.batch_shape.numel()
-            # Sample a single index from a uniform multinomial distribution
-            sampled_idx = torch.multinomial(torch.ones(num_models), 1)
-            # Set weights: one 1 at sampled_idx, rest 0
-            wrapped_model.ensemble_weights = torch.zeros(num_models, 1)
-            none_throws(wrapped_model.ensemble_weights)[sampled_idx] = 1.0
-        else:
-            wrapped_model.ensemble_weights = None
-        surrogate_model._model = wrapped_model
-
     @property
     def surrogate(self) -> TorchAdapter:
         if self._surrogate is None:
             self._surrogate = none_throws(self.get_surrogate)()
-        if not isinstance(
-            # pyre-ignore[16]: `ax.generators.torch_base.TorchGenerator` has no
-            # attribute `surrogate`.
-            self._surrogate.generator.surrogate.model,
-            DeterministicModel,
-        ):
-            self.wrap_surrogate_in_deterministic_model()
         return none_throws(self._surrogate)
 
     def evaluate_true(self, params: Mapping[str, TParamValue]) -> Tensor:
diff --git a/ax/benchmark/testing/benchmark_stubs.py b/ax/benchmark/testing/benchmark_stubs.py
@@ -46,12 +46,10 @@
 from ax.generation_strategy.external_generation_node import ExternalGenerationNode
 from ax.generation_strategy.generation_strategy import GenerationStrategy
 from ax.generators.torch.botorch_modular.generator import BoTorchGenerator
-from ax.generators.torch.botorch_modular.surrogate import ModelConfig, SurrogateSpec
 from ax.utils.testing.core_stubs import (
     get_branin_experiment,
     get_branin_experiment_with_multi_objective,
 )
-from botorch.models.fully_bayesian import SaasFullyBayesianSingleTaskGP
 from botorch.test_functions.multi_objective import BraninCurrin
 from botorch.test_functions.synthetic import Branin
 
@@ -373,33 +371,3 @@ def get_mock_lcbench_data() -> LCBenchData:
         metric_series=metric_series,
         timestamp_series=timestamp_series,
     )
-
-
-def get_adapter(experiment: Experiment) -> TorchAdapter:
-    """Create a generic adapter for testing different surrogate model types."""
-    adapter = TorchAdapter(
-        experiment=experiment,
-        generator=BoTorchGenerator(),
-    )
-    return adapter
-
-
-def get_saas_adapter(experiment: Experiment) -> TorchAdapter:
-    """Create an adapter with SaasFullyBayesianSingleTaskGP model."""
-    return TorchAdapter(
-        experiment=experiment,
-        generator=BoTorchGenerator(
-            surrogate_spec=SurrogateSpec(
-                model_configs=[
-                    ModelConfig(
-                        botorch_model_class=SaasFullyBayesianSingleTaskGP,
-                        mll_options={
-                            "warmup_steps": 2,
-                            "num_samples": 4,
-                            "thinning": 1,
-                        },
-                    ),
-                ]
-            ),
-        ),
-    )
diff --git a/ax/benchmark/tests/benchmark_test_functions/test_surrogate_test_function.py b/ax/benchmark/tests/benchmark_test_functions/test_surrogate_test_function.py
@@ -7,24 +7,11 @@
 
 from unittest.mock import MagicMock, patch
 
-import numpy as np
-
 import torch
 from ax.adapter.torch import TorchAdapter
 from ax.benchmark.benchmark_test_functions.surrogate import SurrogateTestFunction
-from ax.benchmark.testing.benchmark_stubs import (
-    get_adapter,
-    get_saas_adapter,
-    get_soo_surrogate_test_function,
-)
-from ax.generators.torch.botorch_modular.generator import BoTorchGenerator
+from ax.benchmark.testing.benchmark_stubs import get_soo_surrogate_test_function
 from ax.utils.common.testutils import TestCase
-from ax.utils.testing.core_stubs import (
-    get_branin_experiment,
-    get_branin_experiment_with_multi_objective,
-)
-from botorch.models.deterministic import PosteriorMeanModel
-from botorch.sampling.pathwise.posterior_samplers import MatheronPathModel
 
 
 class TestSurrogateTestFunction(TestCase):
@@ -45,160 +32,6 @@ def test_surrogate_test_function(self) -> None:
                 self.assertEqual(test_function.name, "test test function")
                 self.assertIs(test_function.surrogate, surrogate)
 
-    def test_equality(self) -> None:
-        def _construct_test_function(name: str) -> SurrogateTestFunction:
-            return SurrogateTestFunction(
-                name=name,
-                _surrogate=MagicMock(),
-                outcome_names=["dummy_metric"],
-            )
-
-        runner_1 = _construct_test_function("test 1")
-        runner_2 = _construct_test_function("test 2")
-        runner_1a = _construct_test_function("test 1")
-        self.assertEqual(runner_1, runner_1a)
-        self.assertNotEqual(runner_1, runner_2)
-        self.assertNotEqual(runner_1, 1)
-        self.assertNotEqual(runner_1, None)
-
-    def test_surrogate_model_types(self) -> None:
-        """Test different surrogate model types: sample and mean."""
-        experiment = get_branin_experiment(with_completed_trial=True)
-
-        for surrogate_model_type in [MatheronPathModel, PosteriorMeanModel]:
-            with self.subTest(surrogate_model_type=surrogate_model_type):
-                adapter = get_adapter(experiment)
-
-                test_function = SurrogateTestFunction(
-                    name=f"test_{surrogate_model_type}_surrogate",
-                    outcome_names=["branin"],
-                    _surrogate=adapter,
-                    surrogate_model_type=surrogate_model_type,
-                    seed=42,
-                )
-
-                # Verify the surrogate type is set correctly
-                self.assertEqual(
-                    test_function.surrogate_model_type, surrogate_model_type
-                )
-                self.assertEqual(test_function.seed, 42)
-
-                # Test evaluation
-                test_params = {"x1": 0.5, "x2": 0.5}
-                result = test_function.evaluate_true(test_params)
-
-                # Ensure result is a tensor
-                self.assertIsInstance(result, torch.Tensor)
-                self.assertEqual(result.dtype, torch.double)
-                self.assertEqual(result.shape, torch.Size([1]))  # One outcome
-
-    def test_surrogate_model_types_with_random_seeds(self) -> None:
-        """Test that different random seeds produce different results for samples."""
-        experiment = get_branin_experiment(with_completed_trial=True)
-        test_params = {"x1": 0.5, "x2": 0.5}
-
-        results = []
-        for seed in [0, 1, 2]:
-            adapter = get_adapter(experiment)
-            test_function = SurrogateTestFunction(
-                name=f"test_sample_surrogate_seed_{seed}",
-                outcome_names=["branin"],
-                _surrogate=adapter,
-                surrogate_model_type=MatheronPathModel,
-                seed=seed,
-            )
-
-            result = test_function.evaluate_true(test_params)
-            results.append(result.item())
-
-        # Different seeds should produce different results for sample type
-        self.assertFalse(
-            all(r == results[0] for r in results[1:]),
-            "Different random seeds should produce different sample results",
-        )
-
-    def test_mean_surrogate_consistency(self) -> None:
-        """Test that mean surrogate type produces consistent results."""
-        experiment = get_branin_experiment(with_completed_trial=True)
-        test_params = {"x1": 0.5, "x2": 0.5}
-
-        results = []
-        # outcomes should be consistent since seed is fixed
-        for i in range(3):
-            adapter = get_adapter(experiment)
-            test_function = SurrogateTestFunction(
-                name=f"test_mean_surrogate_{i}",
-                outcome_names=["branin"],
-                _surrogate=adapter,
-                surrogate_model_type=MatheronPathModel,
-                seed=42,
-            )
-
-            result = test_function.evaluate_true(test_params)
-            results.append(result.item())
-
-        # Mean type should produce consistent results regardless of seed
-        self.assertTrue(np.all(results[0] == np.array(results)))
-
-    def test_surrogate_model_with_multiple_outcomes(self) -> None:
-        """Test surrogate models with multiple outcome names."""
-        experiment = get_branin_experiment_with_multi_objective(
-            with_completed_trial=True
-        )
-        adapter = TorchAdapter(
-            experiment=experiment,
-            search_space=experiment.search_space,
-            generator=BoTorchGenerator(),
-            data=experiment.lookup_data(),
-            transforms=[],
-        )
-
-        for surrogate_model_type in [MatheronPathModel, PosteriorMeanModel]:
-            with self.subTest(surrogate_model_type=surrogate_model_type):
-                test_function = SurrogateTestFunction(
-                    name=f"test_multi_outcome_{surrogate_model_type}",
-                    outcome_names=["branin_a", "branin_b"],
-                    _surrogate=adapter,
-                    surrogate_model_type=surrogate_model_type,
-                )
-                test_params = {"x1": 0.5, "x2": 0.5}
-                result = test_function.evaluate_true(test_params)
-
-                # Should return 2 outcomes
-                self.assertEqual(result.shape, torch.Size([2]))
-
-    def test_saas_surrogate_model(self) -> None:
-        """Test surrogate test function with SaasFullyBayesianSingleTaskGP model."""
-        experiment = get_branin_experiment(with_completed_trial=True)
-
-        # Create adapter with SaasFullyBayesianSingleTaskGP model
-        adapter = get_saas_adapter(experiment)
-
-        for surrogate_model_type in [MatheronPathModel, PosteriorMeanModel]:
-            with self.subTest(surrogate_model_type=surrogate_model_type):
-                test_function = SurrogateTestFunction(
-                    name=f"test_saas_surrogate_{surrogate_model_type}",
-                    outcome_names=["branin"],
-                    _surrogate=adapter,
-                    surrogate_model_type=surrogate_model_type,
-                    seed=123,
-                )
-
-                # Verify the surrogate type is set correctly
-                self.assertEqual(
-                    test_function.surrogate_model_type, surrogate_model_type
-                )
-                self.assertEqual(test_function.seed, 123)
-
-                # Test evaluation
-                test_params = {"x1": 0.5, "x2": 0.5}
-                result = test_function.evaluate_true(test_params)
-
-                # Ensure result is a tensor with correct properties
-                self.assertIsInstance(result, torch.Tensor)
-                self.assertEqual(result.dtype, torch.double)
-                self.assertEqual(result.shape, torch.Size([1]))  # One outcome
-
     def test_lazy_instantiation(self) -> None:
         test_function = get_soo_surrogate_test_function()
 
@@ -222,46 +55,18 @@ def test_instantiation_raises_with_missing_args(self) -> None:
         ):
             SurrogateTestFunction(name="test runner", outcome_names=[])
 
-    def test_ensemble_sampling(self) -> None:
-        """Test that ensemble sampling works correctly."""
-        experiment = get_branin_experiment(with_completed_trial=True)
-        adapter = get_saas_adapter(experiment)  # Creates ensemble model
-
-        # Test with ensemble sampling enabled (default)
-        test_function = SurrogateTestFunction(
-            name="test_ensemble_sampling_enabled",
-            outcome_names=["branin"],
-            _surrogate=adapter,
-            surrogate_model_type=PosteriorMeanModel,
-            sample_from_ensemble=True,
-        )
-
-        # Access surrogate to trigger wrapping
-        surrogate = test_function.surrogate
-        # pyre-ignore[16]: Access base_model through deterministic wrapper
-        wrapped_model = surrogate.generator.surrogate.model
-
-        # Check that exactly one model has weight 1.0 and others have weight 0.0
-        weights = wrapped_model.ensemble_weights
-        self.assertEqual(weights.sum().item(), 1.0)
-        self.assertEqual((weights == 1.0).sum().item(), 1)
-        self.assertEqual((weights == 0.0).sum().item(), len(weights) - 1)
-
-    def test_ensemble_no_sampling(self) -> None:
-        """Test that ensemble weights remain unchanged when sampling is disabled."""
-        experiment = get_branin_experiment(with_completed_trial=True)
-        adapter = get_saas_adapter(experiment)  # Creates ensemble model
-
-        # Test with ensemble sampling disabled
-        test_function = SurrogateTestFunction(
-            name="test_ensemble_sampling_disabled",
-            outcome_names=["branin"],
-            _surrogate=adapter,
-            surrogate_model_type=PosteriorMeanModel,
-            sample_from_ensemble=False,
-        )
+    def test_equality(self) -> None:
+        def _construct_test_function(name: str) -> SurrogateTestFunction:
+            return SurrogateTestFunction(
+                name=name,
+                _surrogate=MagicMock(),
+                outcome_names=["dummy_metric"],
+            )
 
-        # Access surrogate to trigger wrapping
-        surrogate = test_function.surrogate
-        # pyre-ignore[16]: Access base_model through deterministic wrapper
-        self.assertIsNone(surrogate.generator.surrogate.model.ensemble_weights)
+        runner_1 = _construct_test_function("test 1")
+        runner_2 = _construct_test_function("test 2")
+        runner_1a = _construct_test_function("test 1")
+        self.assertEqual(runner_1, runner_1a)
+        self.assertNotEqual(runner_1, runner_2)
+        self.assertNotEqual(runner_1, 1)
+        self.assertNotEqual(runner_1, None)
