Get rid of some easy pyre-ignore (#4974)

Summary:
Pull Request resolved: #4974

See title

Reviewed By: saitcakmak

Differential Revision: D95264008

fbshipit-source-id: d6a3dc2461027b086b3b34985c092350403d79a9

Author: David Eriksson

ax/adapter/adapter_utils.py

@@ -38,7 +38,7 @@
 from ax.core.parameter import ChoiceParameter, Parameter, ParameterType, RangeParameter
 from ax.core.parameter_constraint import ParameterConstraint
 from ax.core.search_space import SearchSpace, SearchSpaceDigest
-from ax.core.types import TBounds, TCandidateMetadata
+from ax.core.types import TBounds, TCandidateMetadata, TNumeric
 from ax.exceptions.core import DataRequiredError, UserInputError
 from ax.generators.torch.botorch_moo_utils import (
     get_weighted_mc_objective_and_objective_thresholds,
@@ -143,8 +143,11 @@ def extract_search_space_digest(
             else:
                 categorical_features.append(i)
             # at this point we can assume that values are numeric due to transforms
-            discrete_choices[i] = p.values  # pyre-ignore [6]
-            bounds.append((min(p.values), max(p.values)))  # pyre-ignore [6]
+            numeric_values: list[TNumeric] = [
+                assert_is_instance_of_tuple(v, (int, float)) for v in p.values
+            ]
+            discrete_choices[i] = numeric_values
+            bounds.append((min(numeric_values), max(numeric_values)))
         elif isinstance(p, RangeParameter):
             if p.log_scale or p.logit_scale:
                 raise UserInputError(
@@ -154,8 +157,7 @@ def extract_search_space_digest(
                 )
             if p.parameter_type == ParameterType.INT:
                 ordinal_features.append(i)
-                d_choices = list(range(int(p.lower), int(p.upper) + 1))
-                # pyre-ignore [6]
+                d_choices: list[TNumeric] = list(range(int(p.lower), int(p.upper) + 1))
                 discrete_choices[i] = d_choices
             bounds.append((p.lower, p.upper))
         else:

ax/adapter/transforms/choice_encode.py

@@ -108,10 +108,12 @@ def transform_search_space(self, search_space: SearchSpace) -> SearchSpace:
                 dependents = None
                 if p.is_hierarchical:
                     # The dependents of hierarchical parameters need to be updated to
-                    # reflect the changes by encoding.
-                    dependents = {
+                    # reflect the changes by encoding. The encoded keys are ints,
+                    # which is a subtype of TParamValue.
+                    encoded_dependents: dict[TParamValue, list[str]] = {
                         encoding[val]: deps for val, deps in p.dependents.items()
                     }
+                    dependents = encoded_dependents
 
                 transformed_parameters[p_name] = ChoiceParameter(
                     name=p_name,
@@ -130,7 +132,7 @@ def transform_search_space(self, search_space: SearchSpace) -> SearchSpace:
                     # Retain the original sort_values if the parameter is not ordered.
                     # Ordered numeric parameters are always sorted.
                     sort_values=p.sort_values if not p.is_ordered else True,
-                    dependents=dependents,  # pyre-ignore[6]
+                    dependents=dependents,
                 )
             else:
                 transformed_parameters[p.name] = p
@@ -153,7 +155,13 @@ def untransform_observation_features(
                 if p_name in obsf.parameters:
                     # Rounding & casting to int in case a floating point value was
                     # generated. This can happen since generation uses float tensors.
-                    pval = int(round(obsf.parameters[p_name]))  # pyre-ignore [6]
+                    # The value can be int (from data) or float (from generation).
+                    param_val = obsf.parameters[p_name]
+                    pval = (
+                        int(round(param_val))
+                        if isinstance(param_val, float)
+                        else assert_is_instance(param_val, int)
+                    )
                     if pval in reverse_transform:
                         obsf.parameters[p_name] = reverse_transform[pval]
         return observation_features

ax/adapter/transforms/int_to_float.py

@@ -99,9 +99,7 @@ def transform_observation_features(
         for obsf in observation_features:
             for p_name in self.transform_parameters:
                 if p_name in obsf.parameters:
-                    # pyre: param is declared to have type `int` but is used
-                    # pyre-fixme[9]: as type `Optional[typing.Union[bool, float, str]]`.
-                    param: int = obsf.parameters[p_name]
+                    param = assert_is_instance(obsf.parameters[p_name], int)
                     obsf.parameters[p_name] = float(param)
         return observation_features
 
@@ -146,9 +144,7 @@ def untransform_observation_features(
             )
             if self.rounding == "strict":
                 for p_name in present_params:
-                    # pyre: param is declared to have type `float` but is used as
-                    # pyre-fixme[9]: type `Optional[typing.Union[bool, float, str]]`.
-                    param: float = obsf.parameters.get(p_name)
+                    param = assert_is_instance(obsf.parameters.get(p_name), float)
                     obsf.parameters[p_name] = int(round(param))  # TODO: T41938776
             else:
                 if self.contains_constrained_integer:
@@ -191,8 +187,8 @@ def untransform_observation_features(
                     # that satisfies the search space bounds, but this candidate may
                     # not satisfy the parameter constraints.
                     for p_name in present_params:
-                        param = obsf.parameters.get(p_name)
-                        obsf.parameters[p_name] = int(round(param))  # pyre-ignore
+                        param = assert_is_instance(obsf.parameters.get(p_name), float)
+                        obsf.parameters[p_name] = int(round(param))
                 else:  # Update observation if rounding was successful
                     for p_name in present_params:
                         obsf.parameters[p_name] = rounded_parameters[p_name]

ax/adapter/transforms/logit.py

@@ -16,6 +16,7 @@
 from ax.core.parameter import ParameterType, RangeParameter
 from ax.core.search_space import SearchSpace
 from ax.generators.types import TConfig
+from pyre_extensions import assert_is_instance
 from scipy.special import expit, logit
 
 if TYPE_CHECKING:
@@ -58,7 +59,7 @@ def transform_observation_features(
         for obsf in observation_features:
             for p_name in self.transform_parameters:
                 if p_name in obsf.parameters:
-                    param: float = obsf.parameters[p_name]  # pyre-ignore [9]
+                    param = assert_is_instance(obsf.parameters[p_name], float)
                     obsf.parameters[p_name] = logit(param).item()
         return observation_features
 
@@ -78,7 +79,7 @@ def untransform_observation_features(
         for obsf in observation_features:
             for p_name in self.transform_parameters:
                 if p_name in obsf.parameters:
-                    param: float = obsf.parameters[p_name]  # pyre-ignore [9]
+                    param = assert_is_instance(obsf.parameters[p_name], float)
                     obsf.parameters[p_name] = expit(param).item()
         return observation_features
 

ax/adapter/transforms/metrics_as_task.py

@@ -16,7 +16,9 @@
 from ax.core.observation import Observation, ObservationData, ObservationFeatures
 from ax.core.parameter import ChoiceParameter, ParameterType
 from ax.core.search_space import SearchSpace
+from ax.core.types import TParamValueList
 from ax.generators.types import TConfig
+from pyre_extensions import assert_is_instance
 
 if TYPE_CHECKING:
     # import as module to make sphinx-autodoc-typehints happy
@@ -58,10 +60,11 @@ def __init__(
         # Use config to specify metric task map
         if "metric_task_map" not in self.config:
             raise ValueError("config must specify metric_task_map")
-        self.metric_task_map: dict[str, list[str]] = self.config[  # pyre-ignore
-            "metric_task_map"
-        ]
-        self.task_values: list[str] = list(self.metric_task_map.keys())
+        self.metric_task_map: dict[str, list[str]] = assert_is_instance(
+            self.config["metric_task_map"], dict
+        )
+        # Explicitly type to match ChoiceParameter.values expected type
+        self.task_values: TParamValueList = list(self.metric_task_map.keys())
         assert "TARGET" not in self.task_values
         self.task_values.append("TARGET")
 
@@ -136,7 +139,7 @@ def transform_search_space(self, search_space: SearchSpace) -> SearchSpace:
         task_param = ChoiceParameter(
             name="METRIC_TASK",
             parameter_type=ParameterType.STRING,
-            values=self.task_values,  # pyre-ignore
+            values=self.task_values,
             is_ordered=False,
             is_task=True,
             sort_values=True,

ax/adapter/transforms/remove_fixed.py

@@ -22,6 +22,7 @@
 )
 from ax.core.search_space import SearchSpace
 from ax.generators.types import TConfig
+from pyre_extensions import assert_is_instance
 
 if TYPE_CHECKING:
     # import as module to make sphinx-autodoc-typehints happy
@@ -59,8 +60,7 @@ def find_adoptable_descendants(
                 continue
             if isinstance(search_space.parameters[child], FixedParameter):
                 lst_adoptable_descendants += find_adoptable_descendants(
-                    # pyre-ignore[6]: It's a fixed parameter for sure.
-                    search_space.parameters[child],
+                    assert_is_instance(search_space.parameters[child], FixedParameter),
                     search_space=search_space,
                 )
             else:
@@ -141,8 +141,9 @@ def transform_search_space(self, search_space: SearchSpace) -> SearchSpace:
                                 search_space.parameters[child], FixedParameter
                             ):
                                 updated_children += find_adoptable_descendants(
-                                    # pyre-ignore[6]: It's a fixed parameter for sure.
-                                    param=search_space.parameters[child],
+                                    param=assert_is_instance(
+                                        search_space.parameters[child], FixedParameter
+                                    ),
                                     search_space=search_space,
                                 )
                             else:

ax/adapter/transforms/trial_as_task.py

@@ -20,7 +20,7 @@
 from ax.generators.types import TConfig
 from ax.utils.common.constants import Keys
 from ax.utils.common.logger import get_logger
-from pyre_extensions import none_throws
+from pyre_extensions import assert_is_instance, none_throws
 
 if TYPE_CHECKING:
     # import as module to make sphinx-autodoc-typehints happy
@@ -86,10 +86,7 @@ def __init__(
             )
         # Get trial level map
         if "trial_level_map" in self.config:
-            # pyre-ignore [9]
-            trial_level_map: dict[str, dict[int | str, int | str]] = self.config[
-                "trial_level_map"
-            ]
+            trial_level_map = assert_is_instance(self.config["trial_level_map"], dict)
             # Validate
             self.trial_level_map: dict[str, dict[int, int | str]] = {}
             for _p_name, level_dict in trial_level_map.items():

ax/benchmark/benchmark_result.py

@@ -18,6 +18,7 @@
 from ax.utils.common.base import Base
 from numpy import nanmean, nanquantile
 from pandas import DataFrame
+from pyre_extensions import none_throws
 from scipy.stats import sem
 
 PERCENTILES = [0.25, 0.5, 0.75]
@@ -188,11 +189,7 @@ def from_benchmark_results(
         num_trials_mean = None
         if all(res.num_trials is not None for res in results):
             num_trials_step_data = zip(
-                *(
-                    # pyre-ignore[16]: already checked for None above
-                    res.num_trials
-                    for res in results
-                )
+                *(none_throws(res.num_trials) for res in results)
             )
             num_trials_mean = [nanmean(step_vals) for step_vals in num_trials_step_data]