Disallow noise_std as list (#4770)

Summary:
Pull Request resolved: #4770

Per discussion in D90596997

Reviewed By: saitcakmak

Differential Revision: D90777058

fbshipit-source-id: d90616bbfe37f2ad59eff39a7ea91164cf06f1af

Author: Carl Hvarfner

ax/benchmark/benchmark_problem.py

@@ -45,9 +45,9 @@ class BenchmarkProblem(Base):
             data. This will be used by a `BenchmarkRunner`.
         noise_std: Describes how noise is added to the output of the
             `test_function`. If a float, IID random normal noise with that
-            standard deviation is added. A list of floats, or a dict whose keys
-            match `test_functions.outcome_names`, sets different noise
-            standard deviations for the different outcomes produced by the
+            standard deviation is added. A dict whose keys match
+            `test_functions.outcome_names` sets different noise standard
+            deviations for the different outcomes produced by the
             `test_function`. This will be used by a `BenchmarkRunner`.
         optimal_value: The best ground-truth objective value, used for scoring
             optimization results on a scale from 0 to 100, where achieving the
@@ -93,7 +93,7 @@ class BenchmarkProblem(Base):
     optimization_config: OptimizationConfig
     num_trials: int
     test_function: BenchmarkTestFunction
-    noise_std: float | Sequence[float] | Mapping[str, float] = 0.0
+    noise_std: float | Mapping[str, float] = 0.0
     optimal_value: float
     baseline_value: float
     worst_feasible_value: float | None = None

ax/benchmark/benchmark_runner.py

@@ -24,7 +24,6 @@
 from ax.runners.simulated_backend import SimulatedBackendRunner
 from ax.utils.common.serialization import TClassDecoderRegistry, TDecoderRegistry
 from ax.utils.testing.backend_simulator import BackendSimulator, BackendSimulatorOptions
-from pyre_extensions import assert_is_instance
 
 
 def _dict_of_arrays_to_df(
@@ -163,7 +162,7 @@ class BenchmarkRunner(Runner):
         test_function: A ``BenchmarkTestFunction`` from which to generate
             deterministic data before adding noise.
         noise_std: The standard deviation of the noise added to the data. Can be
-            a list or dict to be per-metric.
+            a dict to be per-metric.
         step_runtime_function: A function that takes in parameters
             (in ``TParameterization`` format) and returns the runtime of a step.
         max_concurrency: The maximum number of trials that can be running at a
@@ -176,7 +175,7 @@ class BenchmarkRunner(Runner):
     """
 
     test_function: BenchmarkTestFunction
-    noise_std: float | Sequence[float] | Mapping[str, float] = 0.0
+    noise_std: float | Mapping[str, float] = 0.0
     step_runtime_function: TBenchmarkStepRuntimeFunction | None = None
     max_concurrency: int = 1
     force_use_simulated_backend: bool = False
@@ -243,17 +242,11 @@ def get_noise_stds(self) -> dict[str, float]:
         noise_std = self.noise_std
         if isinstance(noise_std, float | int):
             return {name: float(noise_std) for name in self.outcome_names}
-        elif isinstance(noise_std, dict):
-            if not set(noise_std.keys()) == set(self.outcome_names):
-                raise ValueError(
-                    "Noise std must have keys equal to outcome names if given as "
-                    "a dict."
-                )
-            return noise_std
-        # list of floats
-        return dict(
-            zip(self.outcome_names, assert_is_instance(noise_std, list), strict=True)
-        )
+        if not set(noise_std.keys()) == set(self.outcome_names):
+            raise ValueError(
+                "Noise std must have keys equal to outcome names if given as a dict."
+            )
+        return dict(noise_std)
 
     def run(self, trial: BaseTrial) -> dict[str, BenchmarkTrialMetadata]:
         """Run the trial by evaluating its parameterization(s).

ax/benchmark/problems/synthetic/from_botorch.py

@@ -90,7 +90,7 @@ def create_problem_from_botorch(
     *,
     test_problem_class: type[BaseTestProblem] | str,
     test_problem_kwargs: dict[str, Any],
-    noise_std: float | list[float] = 0.0,
+    noise_std: float | dict[str, float] = 0.0,
     num_trials: int,
     baseline_value: float | None = None,
     name: str | None = None,

ax/benchmark/testing/benchmark_stubs.py

@@ -60,7 +60,7 @@ def get_single_objective_benchmark_problem(
     num_trials: int = 4,
     test_problem_kwargs: dict[str, Any] | None = None,
     report_inference_value_as_trace: bool = False,
-    noise_std: float | list[float] = 0.0,
+    noise_std: float | dict[str, float] = 0.0,
     status_quo_params: TParameterization | None = None,
 ) -> BenchmarkProblem:
     return create_problem_from_botorch(

ax/benchmark/tests/problems/synthetic/test_from_botorch.py

@@ -130,7 +130,7 @@ def test_single_objective_from_botorch(self) -> None:
     def _test_constrained_from_botorch(
         self,
         observe_noise_sd: bool,
-        noise_std: float | list[float],
+        noise_std: float | dict[str, float],
         test_problem_class: type[ConstrainedBaseTestProblem],
     ) -> None:
         ax_problem = create_problem_from_botorch(
@@ -174,7 +174,15 @@ def _test_constrained_from_botorch(
     def test_constrained_soo_from_botorch(self) -> None:
         for observe_noise_sd, noise_std in product(
             [False, True],
-            [0.0, 0.1, [0.1, 0.3, 0.4]],
+            [
+                0.0,
+                0.1,
+                {
+                    "ConstrainedGramacy": 0.1,
+                    "constraint_slack_0": 0.3,
+                    "constraint_slack_1": 0.4,
+                },
+            ],
         ):
             with self.subTest(observe_noise_sd=observe_noise_sd, noise_std=noise_std):
                 self._test_constrained_from_botorch(

ax/benchmark/tests/test_benchmark_runner.py

@@ -98,7 +98,7 @@ def test_runner(self) -> None:
             )
             for modified_bounds, noise_std in product(
                 (None, [(0.0, 2.0)] * 6),
-                (0.0, [0.1] * num_outcomes),
+                (0.0, {name: 0.1 for name in outcome_names}),
             )
         ]
         param_based_cases = [
@@ -108,11 +108,14 @@ def test_runner(self) -> None:
                 num_outcomes,
             )
             for num_outcomes in (1, 2)
-            for noise_std in (0.0, [float(i) for i in range(num_outcomes)])
+            for noise_std in (
+                0.0,
+                {f"objective_{i}": float(i) for i in range(num_outcomes)},
+            )
         ]
         surrogate_cases = [
             (get_soo_surrogate_test_function(lazy=False), noise_std, 1)
-            for noise_std in (0.0, 1.0, [0.0], [1.0])
+            for noise_std in (0.0, 1.0, {"branin": 0.0}, {"branin": 1.0})
         ]
         for test_function, noise_std, num_outcomes in (
             botorch_cases + param_based_cases + surrogate_cases
@@ -137,11 +140,8 @@ def test_runner(self) -> None:
             ):
                 self.assertIs(runner.test_function, test_function)
                 self.assertEqual(runner.outcome_names, outcome_names)
-                if isinstance(noise_std, list):
-                    self.assertEqual(
-                        runner.get_noise_stds(),
-                        dict(zip(runner.outcome_names, noise_std)),
-                    )
+                if isinstance(noise_std, dict):
+                    self.assertEqual(runner.get_noise_stds(), noise_std)
                 else:  # float
                     self.assertEqual(
                         runner.get_noise_stds(),
@@ -244,15 +244,17 @@ def test_runner(self) -> None:
                     set(runner.test_function.outcome_names), set(res.keys())
                 )
 
-                for i, df in enumerate(res.values()):
-                    if isinstance(noise_std, list):
-                        self.assertEqual(df["sem"].item(), noise_std[i])
-                        if all(n == 0 for n in noise_std):
-                            self.assertTrue(np.array_equal(df["mean"], Y[i, :]))
+                for outcome_name, df in res.items():
+                    if isinstance(noise_std, dict):
+                        self.assertEqual(df["sem"].item(), noise_std[outcome_name])
+                        if all(n == 0 for n in noise_std.values()):
+                            Y_idx = runner.outcome_names.index(outcome_name)
+                            self.assertTrue(np.array_equal(df["mean"], Y[Y_idx, :]))
                     else:  # float
                         self.assertEqual(df["sem"].item(), noise_std)
                         if noise_std == 0:
-                            self.assertTrue(np.array_equal(df["mean"], Y[i, :]))
+                            Y_idx = runner.outcome_names.index(outcome_name)
+                            self.assertTrue(np.array_equal(df["mean"], Y[Y_idx, :]))
 
             with self.subTest(f"test `poll_trial_status()`, {test_description}"):
                 self.assertEqual(
@@ -321,14 +323,10 @@ def test_get_noise_stds(self) -> None:
             )
             self.assertDictEqual(runner.get_noise_stds(), expected_noise_sd_dict)
 
-        with self.subTest("list noise_std"):
-            runner = BenchmarkRunner(test_function=test_function, noise_std=[1.0])
-            self.assertDictEqual(runner.get_noise_stds(), expected_noise_sd_dict)
-
     def test_heterogeneous_noise(self) -> None:
         outcome_names = ["objective_0", "constraint"]
         noise_dict = {"objective_0": 0.1, "constraint": 0.05}
-        for noise_std in [[0.1, 0.05], noise_dict]:
+        for noise_std in [noise_dict]:
             runner = BenchmarkRunner(
                 test_function=BoTorchTestFunction(
                     botorch_problem=ConstrainedHartmann(dim=6),