Remove pyre-fixme/pyre-ignore from ax/ small directories (#4981)

Summary:
Pull Request resolved: #4981

Remove pyre-fixme/pyre-ignore suppression comments from 17 files across
several directories: plot/, analysis/, generation_strategy/, benchmark/,
metrics/, global_stopping/ (both source and test files).

Key fixes:
- Widen `AxPlotConfig.__new__` data param to `dict[str, Any] | Figure`,
  eliminating 14 pyre-fixme[6] suppressions across ax/plot/
- Use `none_throws()` for Optional unwrapping
- Use `assert_is_instance()` for type narrowing
- Add proper type annotations (`AnalysisCard`, `SklearnDataset`, `TParamValue`)
- Use `float()` wrapping for numpy scalar arithmetic
- Use `npt.NDArray` instead of bare `np.ndarray`
- Use `int()` for cardinality() return values
- Fix `tuple[int]` -> `tuple[int, int, int]` for RGB color tuples
- Annotate `DISCRETE_COLOR_SCALE` global

Reviewed By: dme65

Differential Revision: D95264987

fbshipit-source-id: 3eb128e5552913ac0bb7127f110079ccbc8fb08c

Author: saitcakmak

ax/analysis/healthcheck/tests/test_early_stopping_healthcheck.py

@@ -11,6 +11,7 @@
 import pandas as pd
 from ax.analysis.healthcheck.early_stopping_healthcheck import EarlyStoppingAnalysis
 from ax.analysis.healthcheck.healthcheck_analysis import HealthcheckStatus
+from ax.core.analysis_card import AnalysisCard
 from ax.core.data import Data
 from ax.core.experiment import Experiment
 from ax.core.objective import MultiObjective, Objective
@@ -39,9 +40,8 @@ def setUp(self) -> None:
             early_stopping_strategy=self.early_stopping_strategy
         )
 
-    def _get_df_dict(self, card: object) -> dict[str, str]:
+    def _get_df_dict(self, card: AnalysisCard) -> dict[str, str]:
         """Extract Property -> Value dict from card dataframe."""
-        # pyre-ignore[16]: card has df attribute
         return {row["Property"]: row["Value"] for _, row in card.df.iterrows()}
 
     def _create_map_data(

ax/benchmark/benchmark_result.py

@@ -171,9 +171,7 @@ def from_benchmark_results(
         """
         # Extract average wall times and standard errors thereof
         fit_time, gen_time = (
-            # pyre-fixme [16]: Item `float` of `typing.Union[numpy.ndarray[typing.Any,
-            # typing.Any], float]` has no attribute `item`.
-            [nanmean(Ts).item(), sem(Ts, ddof=1, nan_policy="propagate").item()]
+            [float(nanmean(Ts)), float(sem(Ts, ddof=1, nan_policy="propagate"))]
             for Ts in zip(*((res.fit_time, res.gen_time) for res in results))
         )
 

ax/generation_strategy/best_model_selector.py

@@ -20,8 +20,7 @@
 from ax.utils.common.func_enum import FuncEnum
 from pyre_extensions import none_throws
 
-# pyre-fixme[24]: Generic type `np.ndarray` expects 2 type parameters.
-ARRAYLIKE = np.ndarray | list[float] | list[np.ndarray]
+ARRAYLIKE = npt.NDArray | list[float] | list[npt.NDArray]
 
 
 class BestModelSelector(ABC, Base):

ax/generation_strategy/dispatch_utils.py

@@ -186,9 +186,7 @@ def _suggest_gp_model(
                 all_range_parameters_are_discrete = False
             else:
                 num_param_discrete_values = parameter.cardinality()
-                # pyre-fixme[58]: `*` is not supported for operand types `int` and
-                #  `Union[float, int]`.
-                num_possible_points *= num_param_discrete_values
+                num_possible_points *= int(num_param_discrete_values)
 
         if should_enumerate_param:
             num_enumerated_combinations *= none_throws(num_param_discrete_values)

ax/global_stopping/tests/test_strategies.py

@@ -297,7 +297,9 @@ def test_multi_objective(self) -> None:
         gss = ImprovementGlobalStoppingStrategy(
             min_trials=3, window_size=3, improvement_bar=0.1
         )
-        objectives = exp.optimization_config.objective.objectives  # pyre-ignore
+        objectives = assert_is_instance(
+            none_throws(exp.optimization_config).objective, MultiObjective
+        ).objectives
         custom_objective_thresholds = [
             ObjectiveThreshold(
                 metric=objectives[0].metric,

ax/metrics/sklearn.py

@@ -40,8 +40,7 @@ class SklearnDataset(StrEnum):
 
 
 @lru_cache(maxsize=8)
-# pyre-fixme[2]: Parameter must be annotated.
-def _get_data(dataset) -> dict[str, npt.NDArray]:
+def _get_data(dataset: SklearnDataset) -> dict[str, npt.NDArray]:
     """Return sklearn dataset, loading and caching if necessary."""
     if dataset is SklearnDataset.DIGITS:
         return datasets.load_digits()
@@ -105,8 +104,7 @@ def __init__(
             )
         if model_type is SklearnModelType.NN:
             if regression:
-                # pyre-fixme[4]: Attribute must be annotated.
-                self._model_cls = MLPRegressor
+                self._model_cls: Any = MLPRegressor
             else:
                 self._model_cls = MLPClassifier
         elif model_type is SklearnModelType.RF:

ax/metrics/tests/test_chemistry.py

@@ -19,8 +19,7 @@
 
 
 class DummyEnum(Enum):
-    # pyre-fixme[35]: Target cannot be annotated.
-    DUMMY: str = "dummy"
+    DUMMY = "dummy"
 
 
 class ChemistryMetricTest(TestCase):

ax/metrics/tests/test_noisy_function.py

@@ -8,17 +8,17 @@
 
 import math
 
+from ax.core.types import TParamValue
 from ax.metrics.noisy_function import GenericNoisyFunctionMetric
 from ax.utils.common.testutils import TestCase
 from ax.utils.testing.core_stubs import get_trial
+from pyre_extensions import none_throws
 
 
 class GenericNoisyFunctionMetricTest(TestCase):
     def test_GenericNoisyFunctionMetric(self) -> None:
-        # pyre-fixme[3]: Return type must be annotated.
-        # pyre-fixme[2]: Parameter must be annotated.
-        def f(params):
-            return params["x"] + 1.0
+        def f(params: dict[str, TParamValue]) -> float:
+            return float(params["x"]) + 1.0
 
         # noiseless
         metric = GenericNoisyFunctionMetric(
@@ -29,8 +29,10 @@ def f(params):
         df = metric.fetch_trial_data(trial).unwrap().df
         self.assertEqual(df["arm_name"].tolist(), ["0_0"])
         self.assertEqual(df["metric_name"].tolist(), ["test_metric"])
-        # pyre-fixme[16]: Optional type has no attribute `parameters`.
-        self.assertEqual(df["mean"].tolist(), [trial.arm.parameters["x"] + 1.0])
+        self.assertEqual(
+            df["mean"].tolist(),
+            [float(none_throws(trial.arm).parameters["x"]) + 1.0],
+        )
         self.assertEqual(df["sem"].tolist(), [0.0])
 
         # noisy
@@ -43,12 +45,16 @@ def f(params):
         df = metric.fetch_trial_data(trial).unwrap().df
         self.assertEqual(df["arm_name"].tolist(), ["0_0"])
         self.assertEqual(df["metric_name"].tolist(), ["test_metric"])
-        self.assertNotEqual(df["mean"].tolist(), [trial.arm.parameters["x"] + 1.0])
+        self.assertNotEqual(
+            df["mean"].tolist(),
+            [float(none_throws(trial.arm).parameters["x"]) + 1.0],
+        )
         self.assertEqual(df["sem"].tolist(), [1.0])
         df = metric.fetch_trial_data(trial, noisy=False).unwrap().df
         self.assertEqual(df["arm_name"].tolist(), ["0_0"])
         self.assertEqual(df["metric_name"].tolist(), ["test_metric"])
-        self.assertEqual(df["mean"].tolist(), [trial.arm.parameters["x"] + 1.0])
+        arm = none_throws(trial.arm)
+        self.assertEqual(df["mean"].tolist(), [float(arm.parameters["x"]) + 1.0])
         self.assertEqual(df["sem"].tolist(), [0.0])
 
         # unknown noise level
@@ -61,6 +67,7 @@ def f(params):
         df = metric.fetch_trial_data(trial).unwrap().df
         self.assertEqual(df["arm_name"].tolist(), ["0_0"])
         self.assertEqual(df["metric_name"].tolist(), ["test_metric"])
-        self.assertEqual(df["mean"].tolist(), [trial.arm.parameters["x"] + 1.0])
-        self.assertEqual(df["mean"].tolist(), [trial.arm.parameters["x"] + 1.0])
+        arm = none_throws(trial.arm)
+        self.assertEqual(df["mean"].tolist(), [float(arm.parameters["x"]) + 1.0])
+        self.assertEqual(df["mean"].tolist(), [float(arm.parameters["x"]) + 1.0])
         self.assertTrue(math.isnan(df["sem"].tolist()[0]))

ax/plot/base.py

@@ -13,6 +13,7 @@
 from ax.core.types import TParameterization
 from ax.utils.common.serialization import named_tuple_to_dict
 from plotly import utils
+from plotly.graph_objs import Figure
 
 
 # Constants used for numerous plots
@@ -43,7 +44,9 @@ class _AxPlotConfigBase(NamedTuple):
 class AxPlotConfig(_AxPlotConfigBase):
     """Config for plots"""
 
-    def __new__(cls, data: dict[str, Any], plot_type: enum.Enum) -> "AxPlotConfig":
+    def __new__(
+        cls, data: dict[str, Any] | Figure, plot_type: enum.Enum
+    ) -> "AxPlotConfig":
         # Convert data to json-encodable form (strips out NamedTuple and numpy
         # array). This is a lossy conversion.
         dict_data = json.loads(

ax/plot/color.py

@@ -23,8 +23,7 @@ class COLORS(enum.Enum):
 
 
 # colors to be used for plotting discrete series
-# pyre-fixme[5]: Global expression must be annotated.
-DISCRETE_COLOR_SCALE = [
+DISCRETE_COLOR_SCALE: list[tuple[int, int, int]] = [
     COLORS.STEELBLUE.value,
     COLORS.CORAL.value,
     COLORS.PINK.value,