Remove pyre-fixme/pyre-ignore from ax/utils/ source files

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).

Differential Revision: D95265024

Author: saitcakmak

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(

ax/utils/common/typeutils.py

@@ -66,8 +66,7 @@ def assert_is_instance_dict(
     return new_dict
 
 
-# pyre-fixme[34]: `T` isn't present in the function's parameters.
-def assert_is_instance_of_tuple(val: V, typ: tuple[type[V], ...]) -> T:
+def assert_is_instance_of_tuple(val: V, typ: tuple[type[V], ...]) -> V:
     """
     Asserts that a value is an instance of any type in a tuple of types.
 
@@ -79,13 +78,10 @@ def assert_is_instance_of_tuple(val: V, typ: tuple[type[V], ...]) -> T:
     """
     if not isinstance(val, typ):
         raise TypeError(f"Value was not of any type {typ!r}:\n{val!r}")
-    # pyre-fixme[7]: Expected `T` but got `V`.
     return val
 
 
-# pyre-fixme[24]: Generic type `type` expects 1 type parameter, use `typing.Type` to
-#  avoid runtime subscripting errors.
-def _argparse_type_encoder(arg: Any) -> type:
+def _argparse_type_encoder(arg: Any) -> type[Any]:
     """
     Transforms arguments passed to `optimizer_argparse.__call__`
     at runtime to construct the key used for method lookup as

ax/utils/measurement/synthetic_functions.py

@@ -44,6 +44,9 @@ def __call__(
         """Simplified way to call the synthetic function and pass the argument
         numbers directly, e.g. `branin(2.0, 3.0)`.
         """
+        values: (
+            tuple[int | float | npt.NDArray, ...] | list[int | float | npt.NDArray]
+        ) = args
         if kwargs:
             if self.required_dimensionality:
                 assert (
@@ -57,16 +60,16 @@ def __call__(
                 f"Function {self.name} expected either all anonymous "
                 "arguments or all keyword arguments."
             )
-            args = list(kwargs.values())  # pyre-ignore[9]
-        for x in args:
+            values = list(kwargs.values())
+        for x in values:
             if isinstance(x, np.ndarray):
                 return self.f(X=x)
             assert np.isscalar(x), (
                 f"Expected numerical arguments or numpy arrays, got {type(x)}."
             )
             if isinstance(x, int):
                 x = float(x)
-        return assert_is_instance(self.f(np.array(args)), float)
+        return assert_is_instance(self.f(np.array(values)), float)
 
     def f(self, X: npt.NDArray) -> float | npt.NDArray:
         """Synthetic function implementation.
@@ -227,9 +230,9 @@ class Aug_Hartmann6(Hartmann6):
 
     _required_dimensionality = 7
     _domain: list[tuple[float, float]] = [(0.0, 1.0) for i in range(7)]
-    # pyre-fixme[15]: `_minimums` overrides attribute defined in `Hartmann6`
-    #  inconsistently.
-    _minimums = [(0.20169, 0.150011, 0.476874, 0.275332, 0.311652, 0.6573, 1.0)]
+    _minimums: list[tuple[float, ...]] = [
+        (0.20169, 0.150011, 0.476874, 0.275332, 0.311652, 0.6573, 1.0)
+    ]
     _fmin: float = -3.32237
     _fmax = 0.0