Remove pyre-fixme/pyre-ignore from ax/utils/ source files (#4983)

Summary:

Remove ~125 pyre-fixme/pyre-ignore suppression comments from 21 source files
in ax/utils/ by applying proper type fixes:

- Use `none_throws()` for Optional unwrapping (e.g., grad access, sim_start_time)
- Use `cast()` for GPyTorch model and FixedFeatureModel narrowing
- Fix covariant TypeVar issue in Result.unwrap_or with new TypeVar U
- Add proper annotations to StreamHandler[Any], Callable[..., Any]
- Use helper functions for dynamic GPyTorch attribute access
- Use `np.asarray()` for numpy return type consistency
- Fix scipy stub limitations with positional args
- Use `record.__dict__` for safe attribute assignment

All 100 tests pass (0 failures across 6 test targets).

Reviewed By: dme65

Differential Revision: D95265024

Author: saitcakmak

ax/adapter/adapter_utils.py

@@ -12,7 +12,7 @@
 from collections.abc import Callable, Iterable, Mapping, MutableMapping, Sequence
 from copy import deepcopy
 from logging import Logger
-from typing import Any, SupportsFloat, TYPE_CHECKING
+from typing import Any, cast, SupportsFloat, TYPE_CHECKING
 
 import numpy as np
 import numpy.typing as npt
@@ -135,16 +135,18 @@ def extract_search_space_digest(
         if isinstance(p, ChoiceParameter):
             if p.is_task:
                 task_features.append(i)
-                target_values[i] = assert_is_instance_of_tuple(
-                    p.target_value, (int, float)
+                target_values[i] = cast(
+                    TNumeric,
+                    assert_is_instance_of_tuple(p.target_value, (int, float)),
                 )
             elif p.is_ordered:
                 ordinal_features.append(i)
             else:
                 categorical_features.append(i)
             # at this point we can assume that values are numeric due to transforms
             numeric_values: list[TNumeric] = [
-                assert_is_instance_of_tuple(v, (int, float)) for v in p.values
+                cast(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)))
@@ -164,15 +166,21 @@ def extract_search_space_digest(
             raise ValueError(f"Unknown parameter type {type(p)}")
         if p.is_fidelity:
             fidelity_features.append(i)
-            target_values[i] = assert_is_instance_of_tuple(p.target_value, (int, float))
+            target_values[i] = cast(
+                TNumeric,
+                assert_is_instance_of_tuple(p.target_value, (int, float)),
+            )
 
     if search_space.is_hierarchical:
         hierarchical_dependencies = {}
 
         for p_name, p in search_space.parameters.items():
             if p.is_hierarchical:
                 hierarchical_dependencies[param_names.index(p_name)] = {
-                    assert_is_instance_of_tuple(parent_value, (int, float)): [
+                    cast(
+                        TNumeric,
+                        assert_is_instance_of_tuple(parent_value, (int, float)),
+                    ): [
                         param_names.index(activated_param)
                         for activated_param in activated_params
                     ]

ax/adapter/transforms/log.py

@@ -114,12 +114,13 @@ def transform_search_space(self, search_space: SearchSpace) -> SearchSpace:
                     ]
                     target_value = p.target_value
                     if target_value is not None:
-                        target_value = math.log10(
-                            assert_is_instance_of_tuple(target_value, (float, int))
-                        )
+                        assert_is_instance_of_tuple(target_value, (float, int))
+                        target_value = math.log10(float(target_value))
                     if dependents is not None:
                         dependents = {
-                            math.log10(assert_is_instance_of_tuple(k, (float, int))): v
+                            math.log10(
+                                float(assert_is_instance_of_tuple(k, (float, int)))
+                            ): v
                             for k, v in dependents.items()
                         }
 

ax/adapter/transforms/one_hot.py

@@ -9,7 +9,7 @@
 from __future__ import annotations
 
 from copy import deepcopy
-from typing import TYPE_CHECKING
+from typing import cast, TYPE_CHECKING
 
 import numpy as np
 import pandas as pd
@@ -32,7 +32,6 @@
 from ax.core.types import TParameterization, TParamValue
 from ax.exceptions.core import UnsupportedError
 from ax.generators.types import TConfig
-from ax.utils.common.typeutils import assert_is_instance_of_tuple
 from pyre_extensions import assert_is_instance
 
 if TYPE_CHECKING:
@@ -176,9 +175,7 @@ def transform_search_space(self, search_space: SearchSpace) -> SearchSpace:
                     # If the dependent is not being transformed, keep it as is.
                     new_deps.extend(oh_param_names_for_p.get(dep, [dep]))
                 new_dependents[val] = new_deps
-            assert_is_instance_of_tuple(
-                p, (ChoiceParameter, FixedParameter)
-            ).dependents = new_dependents
+            cast(ChoiceParameter | FixedParameter, p).dependents = new_dependents
 
         return construct_new_search_space(
             search_space=search_space,

ax/utils/common/equality.py

@@ -23,9 +23,7 @@ def equality_typechecker(eq_func: Callable) -> Callable:
     """
 
     # no type annotation for now; breaks sphinx-autodoc-typehints
-    # pyre-fixme[3]: Return type must be annotated.
-    # pyre-fixme[2]: Parameter must be annotated.
-    def _type_safe_equals(self, other):
+    def _type_safe_equals(self: Any, other: Any) -> bool:
         if not isinstance(other, self.__class__):
             return False
         return eq_func(self, other)
@@ -60,7 +58,6 @@ def same_elements(list1: list[Any], list2: list[Any]) -> bool:
     return all(matched)
 
 
-# pyre-fixme[2]: Parameter annotation cannot contain `Any`.
 def is_ax_equal(one_val: Any, other_val: Any) -> bool:
     """Check for equality of two values, handling lists, dicts, dfs, floats,
     dates, and numpy arrays. This method and ``same_elements`` function

ax/utils/common/executils.py

@@ -80,20 +80,15 @@ def retry_on_exception(
             then there is no wait between retries.
     """
 
-    # pyre-fixme[3]: Return type must be annotated.
-    # pyre-fixme[2]: Parameter must be annotated.
-    def func_wrapper(func):
+    def func_wrapper(func: Callable[..., Any]) -> Callable[..., Any]:
         # Depending on whether `func` is async or not, we use a slightly different
         # wrapper; if wrapping an async function, decorator will await it.
         # `async_actual_wrapper` and `actual_wrapper` are almost exactly the same,
         # except that the former is async and awaits the wrapped function.
         if asyncio.iscoroutinefunction(func):
 
             @functools.wraps(func)
-            # pyre-fixme[53]: Captured variable `func` is not annotated.
-            # pyre-fixme[3]: Return type must be annotated.
-            # pyre-fixme[2]: Parameter must be annotated.
-            async def async_actual_wrapper(*args, **kwargs):
+            async def async_actual_wrapper(*args: Any, **kwargs: Any) -> Any:
                 (
                     retry_exceptions,
                     no_retry_exceptions,
@@ -110,8 +105,6 @@ async def async_actual_wrapper(*args, **kwargs):
                         no_retry_exceptions=no_retry_exceptions,
                         retry_exceptions=retry_exceptions,
                         suppress_errors=suppress_errors,
-                        # pyre-fixme[6]: For 4th param expected `Optional[str]` but
-                        #  got `Optional[List[str]]`.
                         check_message_contains=check_message_contains,
                         last_retry=i >= retries - 1,
                         logger=logger,
@@ -130,10 +123,7 @@ async def async_actual_wrapper(*args, **kwargs):
             return async_actual_wrapper
 
         @functools.wraps(func)
-        # pyre-fixme[53]: Captured variable `func` is not annotated.
-        # pyre-fixme[3]: Return type must be annotated.
-        # pyre-fixme[2]: Parameter must be annotated.
-        def actual_wrapper(*args, **kwargs):
+        def actual_wrapper(*args: Any, **kwargs: Any) -> Any:
             (
                 retry_exceptions,
                 no_retry_exceptions,
@@ -150,8 +140,6 @@ def actual_wrapper(*args, **kwargs):
                     no_retry_exceptions=no_retry_exceptions,
                     retry_exceptions=retry_exceptions,
                     suppress_errors=suppress_errors,
-                    # pyre-fixme[6]: For 4th param expected `Optional[str]` but got
-                    #  `Optional[List[str]]`.
                     check_message_contains=check_message_contains,
                     last_retry=i >= retries - 1,
                     logger=logger,
@@ -178,7 +166,7 @@ def handle_exceptions_in_retries(
     no_retry_exceptions: tuple[type[Exception], ...],
     retry_exceptions: tuple[type[Exception], ...],
     suppress_errors: bool,
-    check_message_contains: str | None,
+    check_message_contains: list[str] | None,
     last_retry: bool,
     logger: Logger | None,
     wrap_error_message_in: str | None,

ax/utils/common/func_enum.py

@@ -17,17 +17,13 @@ class FuncEnum(Enum):
     """A base class for all enums with the following structure: string values that
     map to names of functions, which reside in the same module as the enum."""
 
-    # pyre-ignore[3]: Input constructors will be used to make different inputs,
-    # so we need to allow `Any` return type here.
     def __call__(self, **kwargs: Any) -> Any:
         """Defines a method, by which the members of this enum can be called,
         e.g. ``MyFunctions.F(**kwargs)``, which will call the corresponding
         function registered by the name ``F`` in the enum."""
         return self._get_function_for_value()(**kwargs)
 
-    # pyre-ignore[31]: Expression `typing.Callable[([...], typing.Any)]`
-    # is not a valid type.
-    def _get_function_for_value(self) -> Callable[[...], Any]:
+    def _get_function_for_value(self) -> Callable[..., Any]:
         """Retrieve the function in this module, name of which corresponds to the
         value of the enum member."""
         try:

ax/utils/common/logger.py

@@ -65,8 +65,7 @@ def get_logger(
     return logger
 
 
-# pyre-fixme[24]: Generic type `logging.StreamHandler` expects 1 type parameter.
-def build_stream_handler(level: int = DEFAULT_LOG_LEVEL) -> logging.StreamHandler:
+def build_stream_handler(level: int = DEFAULT_LOG_LEVEL) -> logging.StreamHandler[Any]:
     """Build the default stream handler used for most Ax logging. Sets
     default level to INFO, instead of WARNING.
 
@@ -86,8 +85,7 @@ def build_stream_handler(level: int = DEFAULT_LOG_LEVEL) -> logging.StreamHandle
 def build_file_handler(
     filepath: str,
     level: int = DEFAULT_LOG_LEVEL,
-    # pyre-fixme[24]: Generic type `logging.StreamHandler` expects 1 type parameter.
-) -> logging.StreamHandler:
+) -> logging.StreamHandler[Any]:
     """Build a file handle that logs entries to the given file, using the
     same formatting as the stream handler.
 
@@ -216,8 +214,7 @@ def inner(*args: Any, **kwargs: Any) -> T:
 # Uses a permissive level on the logger, instead make each
 # handler as permissive/restrictive as desired
 ROOT_LOGGER.setLevel(logging.DEBUG)
-# pyre-fixme[24]: Generic type `logging.StreamHandler` expects 1 type parameter.
-ROOT_STREAM_HANDLER: logging.StreamHandler = build_stream_handler()
+ROOT_STREAM_HANDLER: logging.StreamHandler[Any] = build_stream_handler()
 ROOT_LOGGER.addHandler(ROOT_STREAM_HANDLER)
 
 

ax/utils/common/result.py

@@ -114,9 +114,7 @@ def unwrap_err(self) -> E:
         pass
 
     @abstractmethod
-    # pyre-ignore[46]: The type variable `Variable[T](covariant)` is covariant and
-    # cannot be a parameter type.
-    def unwrap_or(self, default: T) -> T:
+    def unwrap_or(self, default: U) -> T | U:
         """Returns the contained Ok value or a provided default."""
 
         pass
@@ -183,7 +181,7 @@ def unwrap(self) -> T:
     def unwrap_err(self) -> NoReturn:
         raise RuntimeError(f"Tried to unwrap_err {self}.")
 
-    def unwrap_or(self, default: U) -> T:
+    def unwrap_or(self, default: U) -> T | U:
         return self._value
 
     def unwrap_or_else(self, op: Callable[[E], T]) -> T:
@@ -243,9 +241,7 @@ def unwrap(self) -> NoReturn:
     def unwrap_err(self) -> E:
         return self._value
 
-    # pyre-ignore[46]: The type variable `Variable[T](covariant)` is covariant and
-    # cannot be a parameter type.
-    def unwrap_or(self, default: T) -> T:
+    def unwrap_or(self, default: U) -> T | U:
         return default
 
     def unwrap_or_else(self, op: Callable[[E], T]) -> T:

ax/utils/common/serialization.py

@@ -72,9 +72,7 @@ def serialize_init_args(
     return properties
 
 
-# pyre-fixme[24]: Generic type `type` expects 1 type parameter, use `typing.Type` to
-#  avoid runtime subscripting errors.
-def extract_init_args(args: dict[str, Any], class_: type) -> dict[str, Any]:
+def extract_init_args(args: dict[str, Any], class_: type[Any]) -> dict[str, Any]:
     """Given a dictionary, extract the arguments required for the
     given class's constructor.
     """

ax/utils/common/testutils.py

@@ -20,11 +20,10 @@
 import unittest
 import warnings
 from collections.abc import Callable, Generator
-from contextlib import AbstractContextManager
 from logging import Logger
 from pstats import Stats
 from types import FrameType
-from typing import Any, TypeVar, Union
+from typing import Any, cast, TypeVar, Union
 
 import numpy as np
 import torch
@@ -60,9 +59,7 @@ def _get_tb_lines(tb: types.TracebackType) -> list[tuple[str, int, str]]:
     return res
 
 
-# pyre-fixme[24]: Generic type `unittest.case._AssertRaisesContext` expects 1 type
-#  parameter.
-class _AssertRaisesContextOn(unittest.case._AssertRaisesContext):
+class _AssertRaisesContextOn(unittest.case._AssertRaisesContext[Exception]):
     """
     Attributes:
        lineno: the line number on which the error occurred
@@ -89,12 +86,10 @@ def __init__(
             expected=expected, test_case=test_case, expected_regex=expected_regex
         )
 
-    # pyre-fixme[14]: `__exit__` overrides method defined in `_AssertRaisesContext`
-    #  inconsistently.
     def __exit__(
         self,
-        exc_type: type[Exception] | None,
-        exc_value: Exception | None,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
         tb: types.TracebackType | None,
     ) -> bool:
         """This is called when the context closes. If an exception was raised
@@ -110,10 +105,8 @@ def __exit__(
         self.filename, self.lineno, _ = frames[0]
         lines = [line for _, _, line in frames]
         if self._expected_line is not None and self._expected_line not in lines:
-            # pyre-ignore [16]: ... has no attribute `_raiseFailure`.
-            self._raiseFailure(
-                f"{self._expected_line!r} was not found in the traceback: {lines!r}"
-            )
+            msg = f"{self._expected_line!r} was not found in the traceback: {lines!r}"
+            raise self.test_case.failureException(msg)
 
         return True
 
@@ -420,12 +413,10 @@ def assertRaisesOn(
         exc: type[Exception],
         line: str | None = None,
         regex: str | None = None,
-        # pyre-ignore[24]: Generic type `AbstractContextManager`
-        # expects 2 type parameters, received 1.
-    ) -> AbstractContextManager[None]:
+    ) -> _AssertRaisesContextOn:
         """Assert that an exception is raised on a specific line."""
         context = _AssertRaisesContextOn(exc, self, line, regex)
-        return context.handle("assertRaisesOn", [], {})
+        return cast(_AssertRaisesContextOn, context.handle("assertRaisesOn", [], {}))
 
     def assertDictsAlmostEqual(
         self, a: dict[str, Any], b: dict[str, Any], consider_nans_equal: bool = False
@@ -532,30 +523,21 @@ def ax_long_test(cls, reason: str | None) -> Generator[None, None, None]:
         cls._long_test_active_reason = None
 
     # This list is taken from the python standard library
-    # pyre-fixme[4]: Attribute must be annotated.
-    failUnlessEqual = assertEquals = _deprecate(unittest.TestCase.assertEqual)
-    # pyre-fixme[4]: Attribute must be annotated.
-    failIfEqual = assertNotEquals = _deprecate(unittest.TestCase.assertNotEqual)
-    # pyre-fixme[4]: Attribute must be annotated.
-    failUnlessAlmostEqual = assertAlmostEquals = _deprecate(
-        unittest.TestCase.assertAlmostEqual
-    )
-    # pyre-fixme[4]: Attribute must be annotated.
-    failIfAlmostEqual = assertNotAlmostEquals = _deprecate(
-        unittest.TestCase.assertNotAlmostEqual
-    )
-    # pyre-fixme[4]: Attribute must be annotated.
-    failUnless = assert_ = _deprecate(unittest.TestCase.assertTrue)
-    # pyre-fixme[4]: Attribute must be annotated.
-    failUnlessRaises = _deprecate(unittest.TestCase.assertRaises)
-    # pyre-fixme[4]: Attribute must be annotated.
-    failIf = _deprecate(unittest.TestCase.assertFalse)
-    # pyre-fixme[4]: Attribute must be annotated.
-    assertRaisesRegexp = _deprecate(unittest.TestCase.assertRaisesRegex)
-    # pyre-fixme[4]: Attribute must be annotated.
-    assertRegexpMatches = _deprecate(unittest.TestCase.assertRegex)
-    # pyre-fixme[4]: Attribute must be annotated.
-    assertNotRegexpMatches = _deprecate(unittest.TestCase.assertNotRegex)
+    failUnlessEqual: Callable = _deprecate(unittest.TestCase.assertEqual)
+    assertEquals: Callable = failUnlessEqual
+    failIfEqual: Callable = _deprecate(unittest.TestCase.assertNotEqual)
+    assertNotEquals: Callable = failIfEqual
+    failUnlessAlmostEqual: Callable = _deprecate(unittest.TestCase.assertAlmostEqual)
+    assertAlmostEquals: Callable = failUnlessAlmostEqual
+    failIfAlmostEqual: Callable = _deprecate(unittest.TestCase.assertNotAlmostEqual)
+    assertNotAlmostEquals: Callable = failIfAlmostEqual
+    failUnless: Callable = _deprecate(unittest.TestCase.assertTrue)
+    assert_: Callable = failUnless
+    failUnlessRaises: Callable = _deprecate(unittest.TestCase.assertRaises)
+    failIf: Callable = _deprecate(unittest.TestCase.assertFalse)
+    assertRaisesRegexp: Callable = _deprecate(unittest.TestCase.assertRaisesRegex)
+    assertRegexpMatches: Callable = _deprecate(unittest.TestCase.assertRegex)
+    assertNotRegexpMatches: Callable = _deprecate(unittest.TestCase.assertNotRegex)
 
     # Copied from BoTorch assertAllClose
     def assertAllClose(