From bae241c547e05d228e9c5bc3f547aca61bc9627f Mon Sep 17 00:00:00 2001
From: Elizabeth Santorella <santorella@meta.com>
Date: Fri, 28 Feb 2025 08:55:27 -0800
Subject: [PATCH] Allow int-valued noise_std in BenchmarkRunner (#3439)

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

**Context:**

Yuck! Python typing says an int is acceptable when annotations indicate the need for a float, but `isinstance(3, float)` is False. https://discuss.python.org/t/clarifying-the-float-int-complex-special-case/54018

That was a causing a failure in a function with an `isinstance` check even though there were no complaints from Pyre.

**This PR:**

`isinstance(x, float)` -> `isinstance(x, float | int)`

Reviewed By: Balandat

Differential Revision: D70392395

fbshipit-source-id: c69d6252dbed87e91d42c03c699c386d72daac4e
---
 ax/benchmark/benchmark_runner.py            |  4 ++--
 ax/benchmark/tests/test_benchmark_runner.py | 24 +++++++++++++++++++++
 2 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/ax/benchmark/benchmark_runner.py b/ax/benchmark/benchmark_runner.py
index ad5ac39695b..0295b91f960 100644
--- a/ax/benchmark/benchmark_runner.py
+++ b/ax/benchmark/benchmark_runner.py
@@ -217,8 +217,8 @@ def get_Y_true(self, params: Mapping[str, TParamValue]) -> npt.NDArray:
 
     def get_noise_stds(self) -> dict[str, float]:
         noise_std = self.noise_std
-        if isinstance(noise_std, float):
-            return {name: noise_std for name in self.outcome_names}
+        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(
diff --git a/ax/benchmark/tests/test_benchmark_runner.py b/ax/benchmark/tests/test_benchmark_runner.py
index f45cb180740..22b3562bcf3 100644
--- a/ax/benchmark/tests/test_benchmark_runner.py
+++ b/ax/benchmark/tests/test_benchmark_runner.py
@@ -302,6 +302,30 @@ def test__add_noise(self) -> None:
         self.assertLess(np.abs(z_scores).min(), 2)
         self.assertGreater(z_scores.max(), 0.05)
 
+    def test_get_noise_stds(self) -> None:
+        test_function = BoTorchTestFunction(
+            botorch_problem=Hartmann(dim=6),
+            outcome_names=["objective_0"],
+        )
+        expected_noise_sd_dict = {"objective_0": 1.0}
+        with self.subTest("float noise_std"):
+            runner = BenchmarkRunner(test_function=test_function, noise_std=1.0)
+            self.assertDictEqual(runner.get_noise_stds(), expected_noise_sd_dict)
+
+        with self.subTest("int noise_std"):
+            runner = BenchmarkRunner(test_function=test_function, noise_std=1)
+            self.assertDictEqual(runner.get_noise_stds(), expected_noise_sd_dict)
+
+        with self.subTest("dict noise_std"):
+            runner = BenchmarkRunner(
+                test_function=test_function, noise_std=expected_noise_sd_dict
+            )
+            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}