[BE] typing for decorators - fx/_compatibility (pytorch#134054)

Summary:
X-link: ctrl-labs/src2#33884

X-link: pytorch/executorch#4810

X-link: pytorch/torchrec#2322

Pull Request resolved: pytorch#134054

See pytorch#131429

Test Plan: unit tests pass

Reviewed By: laithsakka

Differential Revision: D61493706

Author: aorenste

torch/_export/__init__.py

@@ -1,4 +1,3 @@
-# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import copy
 import dataclasses

torch/fx/_compatibility.py

@@ -1,14 +1,15 @@
-# mypy: allow-untyped-defs
-from typing import Any, Dict
+from typing import Any, Dict, Callable, TypeVar
 import textwrap
 
 _BACK_COMPAT_OBJECTS : Dict[Any, None] = {}
 _MARKED_WITH_COMPATIBILITY : Dict[Any, None] = {}
 
-def compatibility(is_backward_compatible : bool):
+_T = TypeVar("_T")
+
+def compatibility(is_backward_compatible: bool) -> Callable[[_T], _T]:
     if is_backward_compatible:
 
-        def mark_back_compat(fn):
+        def mark_back_compat(fn: _T) -> _T:
             docstring = textwrap.dedent(getattr(fn, '__doc__', None) or '')
             docstring += """
 .. note::
@@ -22,7 +23,7 @@ def mark_back_compat(fn):
         return mark_back_compat
     else:
 
-        def mark_not_back_compat(fn):
+        def mark_not_back_compat(fn: _T) -> _T:
             docstring = textwrap.dedent(getattr(fn, '__doc__', None) or '')
             docstring += """
 .. warning::

torch/fx/_lazy_graph_module.py

@@ -1,4 +1,3 @@
-# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from contextlib import contextmanager
 

torch/fx/_symbolic_trace.py

@@ -1,4 +1,3 @@
-# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import builtins
 import copy

torch/fx/experimental/proxy_tensor.py

@@ -78,6 +78,7 @@
 
 if TYPE_CHECKING:
     import types
+    from collections.abc import MutableMapping
 
     import sympy
 
@@ -201,9 +202,11 @@ def set_proxy_slot(
     if isinstance(obj, Tensor):
         # We DO want to clobber proxies whenever we run an inplace operation
         # on a tensor, and it affects the metadata on the proxy.
+        assert isinstance(proxy, _ProxyTensor)
         tracer.tensor_tracker[obj] = proxy
     elif isinstance(obj, (_AnyScriptObject)):
         # We DO want to clobber proxies, with a similar rationale as for tensors.
+        assert isinstance(proxy, Proxy)
         tracer.script_object_tracker[obj] = proxy
     else:
         # NB: Never clobber pre-existing proxy.  Although the proxies
@@ -991,11 +994,13 @@ class PythonKeyTracer(Tracer):
 
     def __init__(self) -> None:
         super().__init__(autowrap_modules=())  # type: ignore[arg-type]
-        self.tensor_tracker = WeakTensorKeyDictionary()
+        self.tensor_tracker: MutableMapping[
+            Tensor, _ProxyTensor
+        ] = WeakTensorKeyDictionary()
         self.symnode_tracker = _SymNodeDict()
-        self.script_object_tracker = WeakIdKeyDictionary(
-            dict=None, ref_type=_WeakHashRef
-        )
+        self.script_object_tracker: MutableMapping[
+            _AnyScriptObjectType, Proxy
+        ] = WeakIdKeyDictionary(dict=None, ref_type=_WeakHashRef)
         self.sympy_expr_tracker: Dict[sympy.Symbol, object] = dict()
 
         # Stores the torch function that was called during tracing
@@ -1363,9 +1368,9 @@ def __sym_dispatch__(
 
 
 class _GraphAppendingTracerEx(fx.proxy.GraphAppendingTracer):
-    script_object_tracker: WeakKeyDictionary
+    script_object_tracker: MutableMapping[_AnyScriptObjectType, Proxy]
     symnode_tracker: WeakKeyDictionary
-    tensor_tracker: WeakTensorKeyDictionary
+    tensor_tracker: MutableMapping[Tensor, _ProxyTensor]
     sympy_expr_tracker: Dict[sympy.Symbol, object]
     torch_fn_metadata: Optional[OpOverload]
     torch_fn_counts: Dict[OpOverload, int]

torch/fx/graph.py

@@ -1,4 +1,3 @@
-# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from collections import defaultdict
 from .node import Node, Argument, Target, map_arg, _type_repr, _get_qualified_name

torch/fx/graph_module.py

@@ -1,4 +1,3 @@
-# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import contextlib
 import copy

torch/fx/interpreter.py

@@ -1,4 +1,3 @@
-# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from .graph_module import GraphModule
 from ._lazy_graph_module import _make_graph_module

torch/fx/node.py

@@ -1,4 +1,3 @@
-# mypy: allow-untyped-decorators
 # Nodes represent a definition of a value in our graph of operators.
 from typing import TYPE_CHECKING, Union, Callable, Any, Tuple, List, Optional, Dict, Set
 from ._compatibility import compatibility

torch/fx/operator_schemas.py

@@ -1,4 +1,3 @@
-# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import torch
 import inspect

torch/fx/passes/graph_manipulation.py

@@ -1,4 +1,3 @@
-# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from typing import Any, Dict, List, NamedTuple, Optional
 

torch/fx/passes/infra/pass_manager.py

@@ -1,4 +1,3 @@
-# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import inspect
 import logging

torch/fx/passes/operator_support.py

@@ -1,4 +1,3 @@
-# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import abc
 import typing as t

torch/fx/passes/param_fetch.py

@@ -1,4 +1,3 @@
-# mypy: allow-untyped-decorators
 from torch.fx.graph_module import GraphModule
 from typing import Any, Callable, Dict, List, Tuple, Type
 import torch

torch/fx/passes/runtime_assert.py

@@ -1,4 +1,3 @@
-# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import logging
 import operator

torch/fx/passes/split_module.py

@@ -1,4 +1,3 @@
-# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import inspect
 from typing import Any, Callable, Dict, List, Optional, Set

torch/fx/passes/split_utils.py

@@ -1,4 +1,3 @@
-# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import copy
 from dataclasses import dataclass, field

torch/fx/passes/splitter_base.py

@@ -1,4 +1,3 @@
-# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import argparse
 import copy

torch/fx/passes/tools_common.py

@@ -1,4 +1,3 @@
-# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from typing import List, Tuple, Union, Dict, Any, Set, Mapping, Optional
 import collections

torch/fx/passes/utils/common.py

@@ -1,4 +1,3 @@
-# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from typing import Dict, Tuple
 

torch/fx/passes/utils/fuser_utils.py

@@ -1,4 +1,3 @@
-# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import copy
 from queue import SimpleQueue

torch/fx/subgraph_rewriter.py

@@ -1,4 +1,3 @@
-# mypy: allow-untyped-decorators
 from .graph_module import GraphModule
 from .graph import Graph
 from .node import Node

torch/fx/traceback.py

@@ -1,4 +1,3 @@
-# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import traceback
 from contextlib import contextmanager

torch/nested/_internal/ops.py

@@ -492,7 +492,7 @@ def linear_backward_default(func, *args, **kwargs):
 
 @register_jagged_func(torch.ops.aten.to.dtype, "input: jt_all, dtype: any")
 def to_dtype(func, *args, **kwargs):
-    _, new_kwargs = normalize_function(
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
         func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
     )
 
@@ -809,7 +809,7 @@ def chunk_default(func, *args, **kwargs):
         lengths = inp._offsets.diff()
         chunked_lengths = lengths.chunk(chunks)
         chunked_offsets = [torch.cumsum(x, dim=0) for x in chunked_lengths]
-        chunked_offsets = [F.pad(x, (1, 0), value=0) for x in chunked_offsets]
+        chunked_offsets = [F.pad(x, (1, 0), value=0) for x in chunked_offsets]  # type: ignore[arg-type]
         nested_kwargs = [
             {"offsets": per_offsets, "_ragged_idx": inp._ragged_idx}
             for per_offsets in chunked_offsets
@@ -1613,7 +1613,7 @@ def masked_select_default(func, *args, **kwargs):
             f"Mask with shape {mask.shape} is not compatible with input's shape {inp.shape}"
         )
     res_values = inp._values.masked_select(mask.values())
-    mask_cumsum = F.pad(mask.values().cumsum(dim=0), (1, 0))
+    mask_cumsum = F.pad(mask.values().cumsum(dim=0), (1, 0))  # type: ignore[arg-type]
 
     args = extract_kwargs(inp)
     args["offsets"] = mask_cumsum[inp._offsets]

torch/onnx/_internal/onnxruntime.py

@@ -1,4 +1,3 @@
-# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import dataclasses
 import importlib