[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

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

@@ -38,6 +38,7 @@
 )
 from typing_extensions import Concatenate, ParamSpec, Self
 from weakref import WeakKeyDictionary
+from collections.abc import MutableMapping
 
 import torch
 import torch._ops
@@ -201,6 +202,7 @@ 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.
@@ -991,7 +993,7 @@ 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
@@ -1365,7 +1367,7 @@ def __sym_dispatch__(
 class _GraphAppendingTracerEx(fx.proxy.GraphAppendingTracer):
     script_object_tracker: WeakKeyDictionary
     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/onnx/_internal/onnxruntime.py

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