[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/_dynamo/compiled_autograd.py

@@ -445,7 +445,7 @@ def to_proxy(self, t):
 
     def bind_tensors_to_proxies(self, tensors, proxies):
         if isinstance(proxies, torch.fx.Proxy):
-            proxies = [proxies[i] for i in range(len(tensors))]
+            proxies = [proxies[i] for i in range(len(tensors))]  # type: ignore[index]
         assert len(tensors) == len(proxies)
         track_tensor_tree(tensors, proxies, constant=None, tracer=self.fx_tracer)
 

torch/_dynamo/eval_frame.py

@@ -1539,7 +1539,7 @@ def result_capturing_wrapper(*graph_inputs):
             # Running graph with interpreter is needed for propagating the stack_trace
             def graph_with_interpreter(*args):
                 with torch.fx.traceback.preserve_node_meta():
-                    return torch.fx.Interpreter(graph).run(*args)
+                    return torch.fx.Interpreter(graph).run(*args)  # type: ignore[arg-type]
 
             with unset_fake_temporarily(), enable_python_dispatcher(), fake_mode:
                 try:
@@ -1561,9 +1561,9 @@ def graph_with_interpreter(*args):
 
             assert graph is not None
             for node in graph.graph.find_nodes(op="get_attr"):
-                if isinstance(getattr(graph, node.target), torch.Tensor):
+                if isinstance(getattr(graph, node.target), torch.Tensor):  # type: ignore[arg-type]
                     node.meta["val"] = fake_mode.from_tensor(
-                        getattr(graph, node.target), static_shapes=True
+                        getattr(graph, node.target), static_shapes=True  # type: ignore[arg-type]
                     )
 
         if same_signature:

torch/_dynamo/utils.py

@@ -2151,7 +2151,7 @@ def get_real_value(node, tracer):
         return cache[node]
 
     op = node.op
-    args, kwargs = torch.fx.node.map_arg(
+    args, kwargs = torch.fx.node.map_arg(  # type: ignore[misc]
         (node.args, node.kwargs),
         lambda n: get_real_value(n, tracer),
     )

torch/_export/__init__.py

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

torch/_export/converter.py

@@ -1343,7 +1343,7 @@ def __init__(
             self.name_to_node,
             # Dummy node.
             torch.fx.Node(
-                None,
+                None,  # type: ignore[arg-type]
                 "mock",
                 "call_function",
                 lambda: None,

torch/_export/pass_base.py

@@ -68,7 +68,7 @@ def __init__(self, callback: "_ExportPassBaseDeprecatedDoNotUse", codegen: CodeG
             self.fake_tensor_mode: Optional[FakeTensorMode] = None
             self.submodules: Dict[torch.nn.Module, str] = {}
 
-        def trace(self) -> None:
+        def trace(self) -> None:  # type: ignore[override]
             raise ExportPassBaseError("ExportTracer doesn't support trace().")
 
         def create_arg(self, a: Argument) -> torch.fx.Node:
@@ -160,7 +160,7 @@ def __init__(self, callback: "_ExportPassBaseDeprecatedDoNotUse", gm: fx.GraphMo
 
         def placeholder(
             self,
-            target: str,
+            target: str,  # type: ignore[override]
             args: Tuple[Argument, ...],
             kwargs: Dict[str, Argument],
         ) -> ProxyValue:
@@ -218,7 +218,7 @@ def call_function(
                 raise ExportPassBaseError(f"Unsupported target type: {target}")
 
         def get_attr(
-            self, target: str, args: Tuple[Argument, ...], kwargs: Dict[str, Argument]
+            self, target: str, args: Tuple[Argument, ...], kwargs: Dict[str, Argument]  # type: ignore[override]
         ) -> Argument:
             return super().get_attr(target, args, kwargs)
 
@@ -231,7 +231,7 @@ def call_module(
             raise ExportPassBaseError("call_module is not supported.")
 
         def call_method(
-            self, target: str, args: Tuple[Argument, ...], kwargs: Dict[str, Argument]
+            self, target: str, args: Tuple[Argument, ...], kwargs: Dict[str, Argument]  # type: ignore[override]
         ) -> None:
             raise ExportPassBaseError("call_method is not supported.")
 
@@ -394,7 +394,7 @@ def call_submodule(
         )
         self.tracer.fake_tensor_mode = prev_tracer.fake_tensor_mode
         interpreter = self.ExportInterpreter(self, graph_module)
-        prev_interpreter, self.interpreter = self.interpreter, torch.fx.Interpreter(
+        prev_interpreter, self.interpreter = self.interpreter, torch.fx.Interpreter(  # type: ignore[assignment]
             torch.fx.GraphModule(torch.nn.Module(), torch.fx.Graph())
         )
         inputs_data = pytree.tree_map_only(ProxyValue, lambda x: x.data, inputs)

torch/_export/passes/replace_with_hop_pass_util.py

@@ -64,7 +64,7 @@ def set_hoo_node_meta(call_func_node):
                 # Rename the name of getitem nodes to the actual name of its contents
                 # for passing verifier and better readability, also propagate metadata
                 for get_item_node in call_func_node.users.keys():
-                    idx: int = get_item_node.args[1]
+                    idx: int = get_item_node.args[1]  # type: ignore[assignment]
                     output_node = output_args[idx]
                     get_item_node._rename(output_node.name)
                     get_item_node.meta = output_node.meta

torch/_functorch/partitioners.py

@@ -186,7 +186,7 @@ def _extract_graph_with_inputs_outputs(
             # joint_graph.nodes).
             continue
         elif node.op == "placeholder":
-            env[node] = InvalidNode
+            env[node] = InvalidNode  # type: ignore[assignment]
         elif node.op == "call_function":
             all_args = pytree.arg_tree_leaves(*node.args, **node.kwargs)
             all_args = [
@@ -195,7 +195,7 @@ def _extract_graph_with_inputs_outputs(
                 if isinstance(x, fx.Node)
             ]
             if any(all_args):
-                env[node] = InvalidNode
+                env[node] = InvalidNode  # type: ignore[assignment]
                 continue
             env[node] = new_graph.node_copy(node, lambda x: env[x])
         elif node.op == "get_attr":

torch/_higher_order_ops/flex_attention.py

@@ -817,7 +817,7 @@ def trace_flex_attention_backward(
         )
     assert isinstance(proxy_mode.tracer, torch.fx.Tracer)
     block_mask = block_mask[:-1] + (mask_graph,)
-    proxy_mode.tracer.root.register_module("fw_graph", fw_graph)
+    proxy_mode.tracer.root.register_module("fw_graph", fw_graph)  # type: ignore[arg-type]
     proxy_mode.tracer.root.register_module("joint_graph", joint_graph)
     proxy_mode.tracer.root.register_module("mask_graph", mask_graph)
     node_args = (

torch/_inductor/constant_folding.py

@@ -222,7 +222,7 @@ def insertable_tensor_check(self, tensor: torch.Tensor) -> bool:
     def add_node_replacement(self, node: torch.fx.Node, tensor: torch.Tensor) -> None:
         self.node_replacements[node] = tensor
 
-    def run(self) -> Any:
+    def run(self) -> Any:  # type: ignore[override]
         env: Dict[torch.fx.Node, Any] = {}
         self.insert_placerholder_values(env)
         return super().run(initial_env=env)

torch/_inductor/debug.py

@@ -144,7 +144,7 @@ def func1(*args: Any) -> int:
         kwargs = {}
         if hasattr(snode, "get_device"):
             kwargs = {"device": snode.get_device()}
-        fx_node = graph.call_function(node_func, args=(), kwargs=kwargs)
+        fx_node = graph.call_function(node_func, args=(), kwargs=kwargs)  # type: ignore[arg-type]
 
         def in_output(snode: Union[BaseSchedulerNode, FusedSchedulerNode]) -> bool:
             if isinstance(snode, FusedSchedulerNode):

torch/_inductor/fx_passes/binary_folding.py

@@ -154,17 +154,17 @@ def _check_conv_and_broadcast_op(conv_node, other):
             return False
         if isinstance(other, torch.fx.Node) and other.op == "get_attr":
             other_meta_value = other.meta.get("val")
-            if not other_meta_value.is_floating_point():
+            if not other_meta_value.is_floating_point():  # type: ignore[union-attr]
                 return False
             if (
-                torch.promote_types(other_meta_value.dtype, weight_meta_value.dtype)
+                torch.promote_types(other_meta_value.dtype, weight_meta_value.dtype)  # type: ignore[union-attr]
                 != weight_meta_value.dtype
             ):
                 if not conv_node.meta.get("_allow_conv_mixed_dtype_folding", False):
                     return False
 
                 if (
-                    other_meta_value.dtype != torch.float
+                    other_meta_value.dtype != torch.float  # type: ignore[union-attr]
                     and weight_meta_value.dtype not in (torch.float16, torch.bfloat16)
                 ):
                     return False

torch/_inductor/fx_passes/efficient_conv_bn_eval.py

@@ -213,7 +213,7 @@ def efficient_conv_bn_eval_graph_transform_inlined(match: Match, *args, **kwargs
         new_node = graph.create_node(
             op="call_function",
             target=efficient_conv_bn_eval_decomposed,
-            args=args,
+            args=args,  # type: ignore[arg-type]
             name="efficient_conv_bn_eval",
         )
 
@@ -223,7 +223,7 @@ def efficient_conv_bn_eval_graph_transform_inlined(match: Match, *args, **kwargs
     # take care of the deletion order:
     # delete bn_node first, and then conv_node
     graph.erase_node(bn_node)
-    graph.erase_node(conv_node)
+    graph.erase_node(conv_node)  # type: ignore[arg-type]
 
     return
 
@@ -304,7 +304,7 @@ def efficient_conv_bn_eval_graph_transform_decomposed(match: Match, *args, **kwa
         new_node = graph.create_node(
             op="call_function",
             target=efficient_conv_bn_eval_decomposed,
-            args=args,
+            args=args,  # type: ignore[arg-type]
             name="efficient_conv_bn_eval",
         )
 
@@ -314,7 +314,7 @@ def efficient_conv_bn_eval_graph_transform_decomposed(match: Match, *args, **kwa
     # take care of the deletion order:
     # delete bn_node first, and then conv_node
     graph.erase_node(bn_node)
-    graph.erase_node(conv_node)
+    graph.erase_node(conv_node)  # type: ignore[arg-type]
 
     return
 
@@ -373,7 +373,7 @@ def efficient_conv_bn_eval_graph_transform(match: Match, *args, **kwargs):
     # Find a pair of conv and bn computation nodes to optimize.
     counters["inductor"]["efficient_conv_bn_eval"] += 1
 
-    with graph.inserting_before(conv_node):
+    with graph.inserting_before(conv_node):  # type: ignore[arg-type]
         # create `get_attr` node to access modules
         # note that we directly call `create_node` to fill the `name`
         # argument. `graph.get_attr` and
@@ -403,4 +403,4 @@ def efficient_conv_bn_eval_graph_transform(match: Match, *args, **kwargs):
     # take care of the deletion order:
     # delete bn_node first, and then conv_node
     graph.erase_node(bn_node)
-    graph.erase_node(conv_node)
+    graph.erase_node(conv_node)  # type: ignore[arg-type]

torch/_inductor/fx_passes/freezing_patterns.py

@@ -223,5 +223,5 @@ def unnecessary_dtype_convert(match: Match, **kwargs):
     """Remove unnecessary dtype conversion op, probably left as a result of Conv-Bn folding"""
     graph = match.graph
     node = match.output_node()
-    node.replace_all_uses_with(node.args[0])
+    node.replace_all_uses_with(node.args[0])  # type: ignore[arg-type]
     graph.erase_node(node)

torch/_inductor/fx_passes/joint_graph.py

@@ -511,7 +511,7 @@ def pointless_view(match: Match, arg, size):
     node = match.output_node()
     arg_size = list(node.args[0].meta["val"].shape)  # type: ignore[union-attr]
     if size == arg_size:
-        node.replace_all_uses_with(node.args[0])
+        node.replace_all_uses_with(node.args[0])  # type: ignore[arg-type]
         match.erase_nodes(graph)
 
 

torch/_inductor/fx_passes/post_grad.py

@@ -1033,7 +1033,7 @@ def is_index_put_and_requires_h2d_sync_for_gpu_value(node):
     # if the value we are putting is a cpu scalar.
     # Therefore, when inductor sees an index_put_ with byte tensor indices,
     # it should *not* convert the cpu scalar value into a gpu tensor.
-    args_, kwargs_ = normalize_function(node.target, node.args, node.kwargs)
+    args_, kwargs_ = normalize_function(node.target, node.args, node.kwargs)  # type: ignore[misc]
     any_byte_bool_indices = False
     indices = args_[1]
     for i in indices:

torch/_inductor/fx_passes/quantization.py

@@ -1887,14 +1887,14 @@ def qconv_weight_prepack(match: Match, *args, **kwargs):
             graph.erase_node(conv_node)
             # Erase the dequant pattern
             if dtype == torch.bfloat16:
-                graph.erase_node(convert_to_bf16)  # type: ignore[possibly-undefined]
-            graph.erase_node(dequant_node)
+                graph.erase_node(convert_to_bf16)  # type: ignore[possibly-undefined, arg-type]
+            graph.erase_node(dequant_node)  # type: ignore[arg-type]
             # Erase the dequant per channel pattern
             if clone_node is not None:
-                graph.erase_node(clone_node)
+                graph.erase_node(clone_node)  # type: ignore[arg-type]
             if dtype == torch.bfloat16:
-                graph.erase_node(weight_to_bf16_node)  # type: ignore[possibly-undefined]
-            graph.erase_node(dequant_per_channel)
+                graph.erase_node(weight_to_bf16_node)  # type: ignore[possibly-undefined, arg-type]
+            graph.erase_node(dequant_per_channel)  # type: ignore[arg-type]
             counters["inductor"]["qconv2d_weight_prepack_matcher_count"] += 1
             counters["inductor"]["qconv2d_weight_prepack_matcher_nodes"] += len(
                 match.nodes
@@ -2581,8 +2581,8 @@ def maybe_replace_node(n: torch.fx.Node) -> torch.fx.Node:
 
                     new_args = map_arg(new_quant_node.args, maybe_replace_node)
                     new_kwargs = map_arg(new_quant_node.kwargs, maybe_replace_node)
-                    new_quant_node.args = new_args
-                    new_quant_node.kwargs = new_kwargs
+                    new_quant_node.args = new_args  # type: ignore[assignment]
+                    new_quant_node.kwargs = new_kwargs  # type: ignore[assignment]
                     graph_module.graph.erase_node(quant_node)
 
     graph_module.graph.lint()

torch/_inductor/fx_passes/reinplace.py

@@ -293,7 +293,7 @@ def can_fuse():
             return
 
         src_inp, src_src, src_scatter_view_op = src.args  # type: ignore[union-attr]
-        with graph.inserting_before(src):
+        with graph.inserting_before(src):  # type: ignore[arg-type]
             new_node = graph_call_function(
                 graph,
                 _generalized_scatter,
@@ -316,7 +316,7 @@ def can_fuse():
                 handle_views(new_src)
                 src.replace_all_uses_with(new_src)  # type: ignore[union-attr]
 
-            graph.erase_node(src)
+            graph.erase_node(src)  # type: ignore[arg-type]
 
     for node in graph.nodes:
         if _is_view_op(node.target):

torch/_inductor/fx_passes/split_cat.py

@@ -193,7 +193,7 @@ def normalize_split_base(
         new_split_node = graph.call_function(
             torch.split,
             args=new_args,
-            kwargs=new_kwargs,
+            kwargs=new_kwargs,  # type: ignore[arg-type]
         )
     split_node.replace_all_uses_with(new_split_node)
     new_split_node.meta.update(split_node.meta)
@@ -378,7 +378,7 @@ def normalize_stack_default(match: Match, *args, **kwargs):
 
     with graph.inserting_after(node):
         new_node = graph.call_function(
-            node.target,
+            node.target,  # type: ignore[arg-type]
             args=(tensors,),
             kwargs={"dim": dim},
         )
@@ -529,7 +529,7 @@ def merge_splits(
 
     to_remove = []
 
-    with graph.inserting_before(first_split):
+    with graph.inserting_before(first_split):  # type: ignore[arg-type]
         # Add the new split node
         new_split = graph.call_function(
             torch.split,
@@ -1535,7 +1535,7 @@ def mutate_cat_node(match: Match, split_sections: List[int], dim: int):
             # case 1: the cat uses all getitems from the split
             if len(split_sections) == len(cat_user.args[0]):  # type: ignore[arg-type]
                 # replace the users of the cat node to be the input of the split node
-                cat_user.replace_all_uses_with(split_node.args[0])
+                cat_user.replace_all_uses_with(split_node.args[0])  # type: ignore[arg-type]
                 # remove the cat node
                 graph.erase_node(cat_user)
                 counters["inductor"]["mutate_cat_pass"] += 1
@@ -1947,7 +1947,7 @@ def remove_split_unbind_children(graph: torch.fx.Graph, inputs: List[torch.fx.No
     # check the split node to remove if it has no users
     for node in nodes:
         if len(node.users.keys()) == 0:  # type: ignore[union-attr]
-            graph.erase_node(node)
+            graph.erase_node(node)  # type: ignore[arg-type]
 
 
 # ############pattern to be optimized is#########

torch/_inductor/graph.py

@@ -766,7 +766,7 @@ def get_numel(self, buffer_name: str) -> Union[int, Expr]:
             return self.graph_inputs[buffer_name].get_numel()
         raise KeyError(f"could not find {buffer_name}")
 
-    def run(self, *args: Any) -> Any:
+    def run(self, *args: Any) -> Any:  # type: ignore[override]
         with dynamo_timed("GraphLowering.run"):
             return super().run(*args)
 
@@ -911,9 +911,9 @@ def constant_name(self, name: str, device_override: Optional[torch.device]) -> s
             )
 
     def placeholder(
-        self, target: str, args: Tuple[object], kwargs: Dict[str, object]
+        self, target: str, args: Tuple[object], kwargs: Dict[str, object]  # type: ignore[override]
     ) -> Union[Expr, TensorBox, None]:
-        example = super().placeholder(target, args, kwargs)
+        example = super().placeholder(target, args, kwargs)  # type: ignore[arg-type]
         self.graph_input_names.append(target)
         if isinstance(example, SymTypes):
             expr = example.node.expr
@@ -968,7 +968,7 @@ def placeholder(
                 self.aligned_inputs.add(target)
         return tensor
 
-    def call_function(self, target: Callable, args: Any, kwargs: Dict[str, Any]) -> Any:  # type: ignore[type-arg]
+    def call_function(self, target: Callable, args: Any, kwargs: Dict[str, Any]) -> Any:  # type: ignore[type-arg, override]
         if target is operator.getitem and isinstance(args[0], (list, tuple, dict)):
             return super().call_function(target, args, kwargs)
 
@@ -1023,7 +1023,7 @@ def can_inline_constant(t: torch.Tensor) -> bool:
         return len(t.shape) == 1 and t.shape[0] <= 8
 
     def get_attr(
-        self, target: str, args: Tuple[()], kwargs: Dict[str, object]
+        self, target: str, args: Tuple[()], kwargs: Dict[str, object]  # type: ignore[override]
     ) -> Union[Constant, TensorBox, ir.Subgraph, TorchBindObject]:
         # this is a constant
         value = getattr_recursive(self.module, target)  # type: ignore[arg-type]
@@ -1063,9 +1063,9 @@ def call_method(self, target: Any, args: Any, kwargs: Any) -> NoReturn:
         raise AssertionError
 
     def output(
-        self, target: str, args: Tuple[object], kwargs: Dict[str, object]
+        self, target: str, args: Tuple[object], kwargs: Dict[str, object]  # type: ignore[override]
     ) -> None:
-        result = super().output(target, args, kwargs)
+        result = super().output(target, args, kwargs)  # type: ignore[arg-type]
         if not isinstance(result, (tuple, list)):
             # nested subgraphs can have singleton outputs
             result = (result,)