Typing for common.py (#160362)

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

Author: Lucaskabela

torch/_dynamo/backends/common.py

@@ -1,5 +1,3 @@
-# mypy: ignore-errors
-
 """
 This module provides common utilities and base classes for TorchDynamo backends.
 
@@ -21,6 +19,9 @@
 import contextlib
 import functools
 import logging
+from collections.abc import Iterable
+from typing import Any, Callable
+from typing_extensions import ParamSpec, TypeVar
 from unittest.mock import patch
 
 import torch
@@ -36,13 +37,18 @@
 
 log = logging.getLogger(__name__)
 
+P = ParamSpec("P")
+R = TypeVar("R")
+
 
 class AotAutograd:
-    def __init__(self, **kwargs) -> None:
+    def __init__(self, **kwargs: Any) -> None:
         self.__name__ = "compiler_fn"
         self.kwargs = kwargs
 
-    def __call__(self, gm: torch.fx.GraphModule, example_inputs, **kwargs):
+    def __call__(
+        self, gm: torch.fx.GraphModule, example_inputs: Iterable[Any], **kwargs: Any
+    ) -> Callable[..., Any]:
         if kwargs:
             log.warning("aot_autograd-based backend ignoring extra kwargs %s", kwargs)
 
@@ -66,16 +72,16 @@ def __call__(self, gm: torch.fx.GraphModule, example_inputs, **kwargs):
             counters["aot_autograd"]["not_ok"] += 1
             return gm
 
-        def wrap_bw_compiler(bw_compiler_fn):
-            def _wrapped_bw_compiler(*args, **kwargs):
+        def wrap_bw_compiler(bw_compiler_fn: Callable[P, R]) -> Callable[..., R]:
+            def _wrapped_bw_compiler(*args: P.args, **kwargs: P.kwargs) -> R:
                 # Note [Wrapping bw_compiler in disable]
                 # The two disables here:
                 # - stop TorchDynamo from trying to compile the bw_compiler function itself
                 # - stop TorchDynamo from trying to compile our the generated backwards pass bw_compiler produces
                 return disable(
                     disable(
                         bw_compiler_fn, reason="do not trace backward compiler function"
-                    )(*args, **kwargs),
+                    )(*args, **kwargs),  # type: ignore[misc]
                     reason="do not trace generated backwards pass",
                 )
 
@@ -99,7 +105,9 @@ def _wrapped_bw_compiler(*args, **kwargs):
         # debug asserts slow down compile time noticeably,
         # So only default them on when the aot_eager backend is used.
         if self.kwargs.get("fw_compiler", None) == nop:
-            patch_config = patch("functorch.compile.config.debug_assert", True)
+            patch_config: contextlib.AbstractContextManager[Any] = patch(
+                "functorch.compile.config.debug_assert", True
+            )
         else:
             patch_config = contextlib.nullcontext()
 
@@ -116,11 +124,11 @@ def _wrapped_bw_compiler(*args, **kwargs):
             raise
 
 
-def aot_autograd(**kwargs) -> AotAutograd:
+def aot_autograd(**kwargs: Any) -> AotAutograd:
     return AotAutograd(**kwargs)
 
 
-def mem_efficient_fusion_kwargs(use_decomps):
+def mem_efficient_fusion_kwargs(use_decomps: bool) -> dict[str, Any]:
     from functorch.compile import (
         default_decompositions,
         min_cut_rematerialization_partition,
@@ -140,28 +148,30 @@ def mem_efficient_fusion_kwargs(use_decomps):
     return kwargs
 
 
-def fake_tensor_unsupported(fn):
+def fake_tensor_unsupported(fn: Callable[[Any, list[Any], Any], R]) -> Any:
     """
     Decorator for backends that need real inputs.  We swap out fake
     tensors for zero tensors.
     """
 
     @functools.wraps(fn)
-    def wrapper(model, inputs, **kwargs):
+    def wrapper(model: Any, inputs: Any, **kwargs: Any) -> Any:
         with _disable_current_modes():
             inputs = list(map(defake, inputs))
-            return fn(model, inputs, **kwargs)
+            return fn(model, inputs, **kwargs)  # type: ignore[call-arg]
 
     return wrapper
 
 
-def device_from_inputs(example_inputs) -> torch.device:
+def device_from_inputs(example_inputs: Iterable[Any]) -> torch.device:
     for x in example_inputs:
         if hasattr(x, "device"):
             return x.device
+    return torch.device("cpu")  # Default fallback
 
 
-def dtype_from_inputs(example_inputs) -> torch.dtype:
+def dtype_from_inputs(example_inputs: Iterable[Any]) -> torch.dtype:
     for x in example_inputs:
         if hasattr(x, "dtype"):
             return x.dtype
+    return torch.float32  # Default fallback