Propagate callable parameter types using ParamSpec (pytorch#142306) (pytorch#151014)

Partially addresses pytorch#142306

Pull Request resolved: pytorch#151014
Approved by: https://github.com/Skylion007

Author: tommyadams5

torch/_dynamo/external_utils.py

@@ -190,11 +190,11 @@ def call_module_hooks_from_backward_state(
 
 
 # used for torch._dynamo.disable(recursive=False)
-def get_nonrecursive_disable_wrapper(fn: Callable[..., Any]) -> Callable[..., Any]:
+def get_nonrecursive_disable_wrapper(fn: Callable[_P, _R]) -> Callable[_P, _R]:
     # wrap function to get the right error message
     # this function is in external_utils so that convert_frame doesn't skip it.
     @functools.wraps(fn)
-    def nonrecursive_disable_wrapper(*args, **kwargs):  # type: ignore[no-untyped-def]
+    def nonrecursive_disable_wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
         return fn(*args, **kwargs)
 
     return nonrecursive_disable_wrapper

torch/_dynamo/testing.py

@@ -465,31 +465,31 @@ def make_test_cls_with_patches(
 
 
 # test Python 3.11+ specific features
-def skipIfNotPy311(fn: Callable[..., Any]) -> Callable[..., Any]:
+def skipIfNotPy311(fn: Callable[_P, _T]) -> Callable[_P, _T]:
     if sys.version_info >= (3, 11):
         return fn
     return unittest.skip(fn)
 
 
-def skipIfNotPy312(fn: Callable[..., Any]) -> Callable[..., Any]:
+def skipIfNotPy312(fn: Callable[_P, _T]) -> Callable[_P, _T]:
     if sys.version_info >= (3, 12):
         return fn
     return unittest.skip("Requires Python 3.12+")(fn)
 
 
-def xfailIfPy312(fn: Callable[..., Any]) -> Callable[..., Any]:
+def xfailIfPy312(fn: Callable[_P, _T]) -> Callable[_P, _T]:
     if sys.version_info >= (3, 12):
         return unittest.expectedFailure(fn)
     return fn
 
 
-def skipIfPy312(fn: Callable[..., Any]) -> Callable[..., Any]:
+def skipIfPy312(fn: Callable[_P, _T]) -> Callable[_P, _T]:
     if sys.version_info >= (3, 12):
         return unittest.skip("Not supported in Python 3.12+")(fn)
     return fn
 
 
-def requiresPy310(fn: Callable[..., Any]) -> Callable[..., Any]:
+def requiresPy310(fn: Callable[_P, _T]) -> Callable[_P, _T]:
     if sys.version_info >= (3, 10):
         return fn
     else:
@@ -498,19 +498,19 @@ def requiresPy310(fn: Callable[..., Any]) -> Callable[..., Any]:
 
 # Controls tests generated in test/inductor/test_torchinductor_dynamic_shapes.py
 # and test/dynamo/test_dynamic_shapes.py
-def expectedFailureDynamic(fn: Callable[..., Any]) -> Callable[..., Any]:
+def expectedFailureDynamic(fn: Callable[_P, _T]) -> Callable[_P, _T]:
     fn._expected_failure_dynamic = True  # type: ignore[attr-defined]
     return fn
 
 
 # Controls tests generated in test/inductor/test_torchinductor_codegen_dynamic_shapes.py
-def expectedFailureCodegenDynamic(fn: Callable[..., Any]) -> Callable[..., Any]:
+def expectedFailureCodegenDynamic(fn: Callable[_P, _T]) -> Callable[_P, _T]:
     fn._expected_failure_codegen_dynamic = True  # type: ignore[attr-defined]
     return fn
 
 
 # Controls test generated in test/inductor/test_cpp_wrapper.py
-def expectedFailureDynamicWrapper(fn: Callable[..., Any]) -> Callable[..., Any]:
+def expectedFailureDynamicWrapper(fn: Callable[_P, _T]) -> Callable[_P, _T]:
     fn._expected_failure_dynamic_wrapper = True  # type: ignore[attr-defined]
     return fn
 

torch/_dynamo/variables/torch.py

@@ -203,7 +203,7 @@ def get_overridable_functions():
     from torch.overrides import get_overridable_functions as get_overridable_functions_
 
     funcs = set(chain.from_iterable(get_overridable_functions_().values()))
-    more = {
+    more: set[Callable[..., Any]] = {
         torch.ones,
         torch.ones_like,
         torch.zeros,

torch/_inductor/utils.py

@@ -1649,8 +1649,8 @@ def save_output_code(code: str) -> None:
 
 
 def run_and_get_kernels(
-    fn: Callable[..., Any], *args: Any, **kwargs: Any
-) -> tuple[Any, list[str]]:
+    fn: Callable[P, _T], *args: P.args, **kwargs: P.kwargs
+) -> tuple[_T, list[str]]:
     result, source_codes = run_and_get_code(fn, *args, **kwargs)
     kernels = []
     for code in source_codes:
@@ -1667,7 +1667,7 @@ def run_with_backward() -> Any:
     return run_and_get_code(run_with_backward)
 
 
-def get_code(fn: Callable[..., Any], *args: Any, **kwargs: Any) -> list[str]:
+def get_code(fn: Callable[P, _T], *args: P.args, **kwargs: P.kwargs) -> list[str]:
     """Get the inductor-generated code, but skip any actual compilation or running."""
     from .graph import GraphLowering
 
@@ -1711,7 +1711,7 @@ def call(self, *args: Any, **kwargs: Any) -> None:
     return source_codes
 
 
-def get_triton_code(fn: Callable[..., Any], *args: Any, **kwargs: Any) -> str:
+def get_triton_code(fn: Callable[P, _T], *args: P.args, **kwargs: P.kwargs) -> str:
     source_codes = get_code(fn, *args, **kwargs)
     # Can have two outputs if backwards was eagerly compiled
     assert 1 <= len(source_codes) <= 2, (
@@ -1720,7 +1720,9 @@ def get_triton_code(fn: Callable[..., Any], *args: Any, **kwargs: Any) -> str:
     return source_codes[0]
 
 
-def run_and_get_triton_code(fn: Callable[..., Any], *args: Any, **kwargs: Any) -> str:
+def run_and_get_triton_code(
+    fn: Callable[P, _T], *args: P.args, **kwargs: P.kwargs
+) -> str:
     _, source_codes = run_and_get_code(fn, *args, **kwargs)
     # Can have two outputs if backwards was eagerly compiled
     assert 1 <= len(source_codes) <= 2, (
@@ -1730,7 +1732,7 @@ def run_and_get_triton_code(fn: Callable[..., Any], *args: Any, **kwargs: Any) -
 
 
 def run_and_get_graph_lowering(
-    fn: Callable[..., Any], *args: Any, **kwargs: Any
+    fn: Callable[P, _T], *args: P.args, **kwargs: P.kwargs
 ) -> tuple[Any, list[GraphLowering]]:
     from torch._inductor.graph import GraphLowering
     from torch._inductor.output_code import CompiledFxGraph
@@ -2386,8 +2388,8 @@ def maybe_get_suppress_shape_guards_ctx() -> contextlib.AbstractContextManager[N
 
 
 def run_and_get_cpp_code(
-    fn: Callable[..., Any], *args: Any, **kwargs: Any
-) -> tuple[Any, str]:
+    fn: Callable[P, _T], *args: P.args, **kwargs: P.kwargs
+) -> tuple[_T, str]:
     # We use the patch context manager instead of using it as a decorator.
     # In this way, we can ensure that the attribute is patched and unpatched correctly
     # even if this run_and_get_cpp_code function is called multiple times.
@@ -2431,9 +2433,9 @@ def shape_env_from_inputs(inputs: Sequence[InputType]) -> Optional[ShapeEnv]:
 
 
 def align_inputs_from_check_idxs(
-    model: Callable[[list[InputType]], Any],
+    model: Callable[[list[InputType]], _T],
     inputs_to_check: Sequence[int],
-) -> Callable[[list[InputType]], Any]:
+) -> Callable[[list[InputType]], _T]:
     if len(inputs_to_check) == 0:
         return model
 

torch/_prims_common/wrappers.py

@@ -5,7 +5,7 @@
 from collections.abc import Sequence
 from functools import wraps
 from types import GenericAlias
-from typing import Callable, NamedTuple, Optional, overload, TypeVar
+from typing import Callable, NamedTuple, Optional, overload, TypeVar, Union
 from typing_extensions import ParamSpec
 
 import torch
@@ -285,7 +285,8 @@ def _out_wrapper(fn: Callable[_P, _T]) -> Callable[_P, _T]:
         is_factory_fn = all(p in sig.parameters for p in factory_kwargs)
 
         @wraps(fn)
-        def _fn(*args: _P.args, out=None, **kwargs: _P.kwargs):
+        def _fn(*args: _P.args, **kwargs: _P.kwargs):
+            out = kwargs.pop("out", None)
             if is_factory_fn and out is not None:
                 for k in factory_kwargs:
                     out_attr = getattr(out, k)
@@ -450,7 +451,9 @@ def _autograd_impl(*args, **kwargs):
 # TODO: when tracing this will add torch tensors and not TensorMeta objects
 # to the trace -- we should fix this by adding a tracing context and NumberMeta classes
 # TODO: this wrapper is currently untested
-def elementwise_unary_scalar_wrapper(fn: Callable) -> Callable:
+def elementwise_unary_scalar_wrapper(
+    fn: Callable[_P, _T],
+) -> Callable[_P, Union[_T, NumberType]]:
     """
     Allows unary operators that accept tensors to work with Python numbers.
     """

torch/distributed/elastic/control_plane.py

@@ -23,8 +23,8 @@ def _worker_server(socket_path: str) -> Generator[None, None, None]:
         server.shutdown()
 
 
-@contextmanager
 @record
+@contextmanager
 def worker_main() -> Generator[None, None, None]:
     """
     This is a context manager that wraps your main entry function. This combines

torch/distributed/elastic/multiprocessing/errors/__init__.py

@@ -58,7 +58,8 @@
 from datetime import datetime
 from functools import wraps
 from string import Template
-from typing import Any, Callable, Optional, TypeVar
+from typing import Any, Callable, Optional, TypeVar, Union
+from typing_extensions import ParamSpec
 
 from torch.distributed.elastic.utils.logging import get_logger
 
@@ -82,7 +83,8 @@
 _EMPTY_ERROR_DATA = {"message": "<NONE>"}
 _NOT_AVAILABLE = "<N/A>"
 
-T = TypeVar("T")
+_R = TypeVar("_R")
+_P = ParamSpec("_P")
 
 
 @dataclass
@@ -305,8 +307,8 @@ def _format_failure(
 
 
 def record(
-    fn: Callable[..., T], error_handler: Optional[ErrorHandler] = None
-) -> Callable[..., T]:
+    fn: Callable[_P, _R], error_handler: Optional[ErrorHandler] = None
+) -> Callable[_P, Union[_R, None]]:
     """
     Syntactic sugar to record errors/exceptions that happened in the decorated
     function using the provided ``error_handler``.
@@ -346,9 +348,9 @@ def main():
     if not error_handler:
         error_handler = get_error_handler()
 
-    def wrap(f):
+    def wrap(f: Callable[_P, _R]) -> Callable[_P, Union[_R, None]]:
         @wraps(f)
-        def wrapper(*args, **kwargs):
+        def wrapper(*args: _P.args, **kwargs: _P.kwargs):
             assert error_handler is not None  # assertion for mypy type checker
             error_handler.initialize()
             try:

torch/distributed/elastic/timer/file_based_local_timer.py

@@ -13,7 +13,8 @@
 import sys
 import threading
 import time
-from typing import Callable, Optional
+from typing import Callable, Optional, TypeVar
+from typing_extensions import ParamSpec
 
 from torch.distributed.elastic.timer.api import TimerClient, TimerRequest
 from torch.distributed.elastic.timer.debug_info_logging import (
@@ -22,6 +23,9 @@
 from torch.distributed.elastic.utils.logging import get_logger
 
 
+_P = ParamSpec("_P")
+_R = TypeVar("_R")
+
 __all__ = ["FileTimerClient", "FileTimerRequest", "FileTimerServer"]
 
 logger = get_logger(__name__)
@@ -36,8 +40,8 @@ def _retry(max_retries: int, sleep_time: float) -> Callable:
         sleep_time: float, the time to sleep between retries.
     """
 
-    def wrapper(func: Callable) -> Callable:
-        def wrapper(*args, **kwargs):
+    def wrapper(func: Callable[_P, _R]) -> Callable[_P, _R]:
+        def wrapper(*args: _P.args, **kwargs: _P.kwargs):
             for i in range(max_retries):
                 try:
                     return func(*args, **kwargs)

torch/fx/node.py

@@ -6,6 +6,7 @@
 import types
 from collections.abc import Mapping, Sequence
 from typing import Any, Callable, Optional, TYPE_CHECKING, TypeVar, Union
+from typing_extensions import ParamSpec
 
 import torch
 from torch._C import _fx_map_aggregate, _fx_map_arg, _NodeBase
@@ -58,6 +59,8 @@
     ]
 ]
 ArgumentT = TypeVar("ArgumentT", bound=Argument)
+_P = ParamSpec("_P")
+_R = TypeVar("_R")
 
 _legal_ops = dict.fromkeys(
     [
@@ -102,7 +105,7 @@
 
 
 @compatibility(is_backward_compatible=False)
-def has_side_effect(fn: Callable) -> Callable:
+def has_side_effect(fn: Callable[_P, _R]) -> Callable[_P, _R]:
     _side_effectful_functions.add(fn)
     return fn
 

torch/onnx/_internal/exporter/_isolated.py

@@ -1,18 +1,24 @@
 """Isolated calls to methods that may segfault."""
 
-# mypy: allow-untyped-defs
 from __future__ import annotations
 
 import multiprocessing
 import os
 import warnings
-from typing import Callable
+from typing import Any, Callable, TypeVar, TypeVarTuple, Union, Unpack
+from typing_extensions import ParamSpec
 
 
+_P = ParamSpec("_P")
+_R = TypeVar("_R")
+_Ts = TypeVarTuple("_Ts")
+
 _IS_WINDOWS = os.name == "nt"
 
 
-def _call_function_and_return_exception(func, args, kwargs):
+def _call_function_and_return_exception(
+    func: Callable[[Unpack[_Ts]], _R], args: tuple[Unpack[_Ts]], kwargs: dict[str, Any]
+) -> Union[_R, Exception]:
     """Call function and return a exception if there is one."""
 
     try:
@@ -21,7 +27,7 @@ def _call_function_and_return_exception(func, args, kwargs):
         return e
 
 
-def safe_call(func: Callable, *args, **kwargs):
+def safe_call(func: Callable[_P, _R], *args: _P.args, **kwargs: _P.kwargs) -> _R:
     """Call a function in a separate process.
 
     Args:

torch/onnx/_internal/registration.py

@@ -4,6 +4,7 @@
 import warnings
 from collections.abc import Collection, Sequence
 from typing import Callable, Generic, Optional, TypeVar, Union
+from typing_extensions import ParamSpec
 
 from torch.onnx import _constants, errors
 
@@ -51,6 +52,8 @@ def _dispatch_opset_version(
 
 _K = TypeVar("_K")
 _V = TypeVar("_V")
+_R = TypeVar("_R")
+_P = ParamSpec("_P")
 
 
 class OverrideDict(Collection[_K], Generic[_K, _V]):
@@ -287,7 +290,7 @@ def symbolic_b(g: _C.Graph, x: _C.Value, y: _C.Value, arg1: bool) -> _C.Value: .
         ValueError: If the separator '::' is not in the name.
     """
 
-    def wrapper(func: Callable) -> Callable:
+    def wrapper(func: Callable[_P, _R]) -> Callable[_P, _R]:
         decorated = func
         if decorate is not None:
             for decorate_func in decorate:

torch/overrides.py

@@ -29,7 +29,8 @@
 import warnings
 from collections.abc import Iterable
 from functools import wraps
-from typing import Any, Callable, Optional
+from typing import Any, Callable, Optional, TypeVar
+from typing_extensions import ParamSpec
 
 import torch
 from torch._C import (
@@ -58,12 +59,15 @@
     "enable_reentrant_dispatch",
 ]
 
+_P = ParamSpec("_P")
+_R = TypeVar("_R")
+
 
 def _disable_user_warnings(
-    func: Callable,
+    func: Callable[_P, _R],
     regex: str = ".*is deprecated, please use.*",
     module: str = "torch",
-) -> Callable:
+) -> Callable[_P, _R]:
     """
     Decorator that temporarily disables ``UserWarning``s for the given ``module`` if the warning message matches the
     given ``regex`` pattern.
@@ -84,7 +88,7 @@ def _disable_user_warnings(
     """
 
     @wraps(func)
-    def wrapper(*args, **kwargs):
+    def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
         with warnings.catch_warnings():
             warnings.filterwarnings(
                 "ignore", category=UserWarning, message=regex, module=module