Revert "Arm backend: Fix arg-type MyPy errors (pytorch#15367)"

This reverts commit c66078c.

Author: GregoryComer

backends/arm/_passes/annotate_decomposed_matmul.py

@@ -51,7 +51,7 @@ def _match_partition_to_node(
         raise RuntimeError(f"Cannot find an input node which matches, {node}.")
 
     def call(self, graph_module: GraphModule) -> PassResult:
-        matmul_partitions_map = get_source_partitions(
+        matmul_partitions = get_source_partitions(
             graph_module.graph,
             [
                 torch.matmul,
@@ -60,7 +60,7 @@ def call(self, graph_module: GraphModule) -> PassResult:
             None,
         )
         matmul_partitions = list(
-            itertools.chain.from_iterable(matmul_partitions_map.values())
+            itertools.chain.from_iterable(matmul_partitions.values())
         )
         matmul_targets = {
             exir_ops.edge.aten.bmm.default,
@@ -88,7 +88,7 @@ def call(self, graph_module: GraphModule) -> PassResult:
                         # Create new dq-node before matmul
                         dq_node = create_node(
                             graph=graph_module.graph,
-                            op_target=cast(EdgeOpOverload, input_node.target),
+                            op_target=cast(EdgeOpOverload, input_node.target),  # type: ignore[arg-type]
                         )
                         dq_node.args = (node, *input_node.args[1:])
                         matmul_node.replace_input_with(node, dq_node)
@@ -109,7 +109,7 @@ def call(self, graph_module: GraphModule) -> PassResult:
                     # Create q-node after matmul
                     q_node = create_node(
                         graph=graph_module.graph,
-                        op_target=cast(EdgeOpOverload, partition_output.target),
+                        op_target=cast(EdgeOpOverload, partition_output.target),  # type: ignore[arg-type]
                     )
                     matmul_node.replace_all_uses_with(q_node)
                     q_node.args = (matmul_node, *partition_output.args[1:])

backends/arm/_passes/arm_pass_utils.py

@@ -13,10 +13,8 @@
 import torch
 import torch.fx
 from executorch.backends.arm.common.debug import get_node_debug_info
-from executorch.backends.arm.common.type import ensure_type
 from executorch.exir import ExportedProgram
 from executorch.exir.dialects._ops import ops as exir_ops
-from executorch.exir.dialects.edge._ops import EdgeOpOverload
 
 from torch._export.utils import (
     get_buffer,
@@ -83,18 +81,17 @@ def get_param_tensor(
     elif is_lifted_tensor_constant(exp_prog, node):
         return get_lifted_tensor_constant(exp_prog, node)
     elif is_get_attr_node(node):
-        target_node = ensure_type(str, node.target)
         # This is a hack to support both lifted and unlifted graph
         try:
-            return getattr(node.graph.owning_module, target_node)
+            return getattr(node.graph.owning_module, node.target)  # type: ignore[arg-type]
         except AttributeError:
-            return getattr(exp_prog.graph_module, target_node)
+            return getattr(exp_prog.graph_module, node.target)  # type: ignore[arg-type]
     raise RuntimeError(f"unsupported param type, {node.op}.")
 
 
 def create_node(
     graph: torch.fx.Graph,
-    op_target: OpOverload | EdgeOpOverload,
+    op_target: OpOverload,
     args: tuple = (),
     kwargs: Optional[dict] = None,
     quantize: bool = False,

backends/arm/_passes/scalars_to_attribute_pass.py

@@ -49,15 +49,15 @@ def call(self, graph_module: GraphModule) -> PassResult:
                     shape = get_first_fake_tensor(arg).shape
                     biggest_rank = max(biggest_rank, len(shape))
 
-            new_args: list[Node | int] = []
+            new_args = []
             for arg in n.args:
                 if isinstance(arg, Node):
                     new_args.append(arg)
                     continue
                 if isinstance(arg, int) and not torch.is_floating_point(
                     get_first_fake_tensor(n)
                 ):
-                    new_args.append(arg)
+                    new_args.append(arg)  # type: ignore[arg-type]
                     continue
 
                 prefix = "_tensor_constant_"

backends/arm/_passes/to_tosa_memory_format_pass.py

@@ -259,19 +259,13 @@ def insert_tosa_transposes(self, graph_module: torch.fx.GraphModule):
 
         # Transpose outputs if they are in (N)NCHW format
         outputs = output_node.args[0]
-        if not isinstance(outputs, (list, tuple)):
-            raise TypeError(
-                f"Expected output node args to be a list or tuple, got {type(outputs)}"
-            )
         output_dim_orders = output_node.meta.get("original_dim_orders")
         if output_dim_orders is None:
             raise RuntimeError(
                 f"{AnnotateDecomposedMatmulPass.__name__} is required to run at the beginning of the pass pipeline when using {ToTosaMemoryFormatPass.__name__}."
             )
 
-        for output_node_input, output_dim_order in zip(
-            outputs, output_dim_orders, strict=True
-        ):
+        for output_node_input, output_dim_order in zip(outputs, output_dim_orders):  # type: ignore[arg-type]
             if output_dim_order in (
                 NCHW_ORDER,
                 NNCHW_ORDER,

backends/arm/common/type.py

@@ -14,7 +14,6 @@
 import torch
 import torch.fx as fx
 from executorch.backends.arm._passes.arm_pass_utils import get_first_fake_tensor
-from executorch.backends.arm.common.type import ensure_type
 from executorch.backends.arm.operator_support.tosa_supported_operators import (
     register_tosa_support_check,
     SupportedTOSAOperatorCheck,
@@ -138,8 +137,7 @@ def is_node_tosa_supported(
                 return False
 
             # Usage 1 guard
-            index = ensure_type(torch.fx.Node, index)
-            fake_tensor = get_first_fake_tensor(index)
+            fake_tensor = get_first_fake_tensor(index)  # type: ignore[arg-type]
             if len(fake_tensor.size()) > 3:
                 self.reporter.report_reject(
                     node,
@@ -148,8 +146,7 @@ def is_node_tosa_supported(
                 return False
 
         # Usage 3 guard
-        input_node = ensure_type(torch.fx.Node, node.args[0])
-        total_vals = math.prod(get_first_fake_tensor(input_node).shape)
+        total_vals = math.prod(get_first_fake_tensor(node.args[0]).shape)  # type: ignore[arg-type]
         if total_vals > torch.iinfo(torch.int32).max:
             self.reporter.report_reject(
                 node,

backends/arm/operator_support/index_tensor_support.py

@@ -219,7 +219,7 @@ def _is_matmul_node_supported(
         """
         for graph_module in submodules.values():
             graph_module = typing.cast(fx.GraphModule, graph_module)
-            matmul_partitions_map = get_source_partitions(
+            matmul_partitions = get_source_partitions(
                 graph_module.graph,
                 [
                     torch.matmul,
@@ -228,7 +228,7 @@ def _is_matmul_node_supported(
                 None,
             )
             matmul_partitions = list(
-                itertools.chain.from_iterable(matmul_partitions_map.values())
+                itertools.chain.from_iterable(matmul_partitions.values())
             )
             matched_partition = None
             for partition in matmul_partitions:
@@ -406,7 +406,9 @@ def is_node_supported(
                 if input_node.target in ComputeConstantOpsAOT.targeted_ops:
                     # This is not perfect since the input_node can still be rejected by other checks but
                     # this should cover the majority of cases.
-                    if self.is_node_supported({}, input_node):
+                    if self.is_node_supported(
+                        None, input_node  # type: ignore[arg-type] #(we don't use 'submodules')
+                    ):
                         continue
             self.reporter.report_reject(
                 node, f"Non-constant int64 input {input_node.name}"

backends/arm/operator_support/tosa_supported_operators.py

@@ -374,7 +374,7 @@ def transform_for_annotation(self, model: GraphModule) -> GraphModule:
         # TODO: Fix the need to lazily import this.
         from executorch.backends.arm._passes import ArmPassManager
 
-        return ArmPassManager(self.tosa_spec).transform_for_annotation_pipeline(
+        return ArmPassManager(self.tosa_spec).transform_for_annotation_pipeline(  # type: ignore[arg-type]
             graph_module=model
         )
 

backends/arm/quantizer/arm_quantizer.py

@@ -12,7 +12,6 @@
 import torch.fx
 import torch.nn.functional as F
 from executorch.backends.arm.common.debug import get_node_debug_info
-from executorch.backends.arm.common.type import ensure_type
 from executorch.backends.arm.quantizer import QuantizationConfig
 from torch._subclasses import FakeTensor
 
@@ -511,8 +510,7 @@ def any_or_hardtanh_min_zero(n: Node):
         torch.ops.aten.minimum.default,
         torch.ops.aten.maximum.default,
     ):
-        lhs_node = ensure_type(Node, node.args[0])
-        shared_qspec = SharedQuantizationSpec((lhs_node, node))
+        shared_qspec = SharedQuantizationSpec((node.args[0], node))  # type: ignore[arg-type]
         quant_properties.quant_inputs = [
             _QuantProperty(0, input_act_qspec),
             _QuantProperty(
@@ -522,24 +520,22 @@ def any_or_hardtanh_min_zero(n: Node):
         ]
         quant_properties.quant_output = _QuantProperty(0, shared_qspec)
     elif node.target in (torch.ops.aten.where.self,):
-        true_node = ensure_type(Node, node.args[1])
-        shared_qspec = SharedQuantizationSpec(true_node)
+        shared_qspec = SharedQuantizationSpec(node.args[1])  # type: ignore[arg-type]
         quant_properties.quant_inputs = [
             _QuantProperty(1, shared_qspec),
             _QuantProperty(2, shared_qspec),
         ]
         quant_properties.quant_output = _QuantProperty(0, shared_qspec)
     elif node.target in _one_to_one_shared_input_or_input_act_qspec:
-        input_node = ensure_type(Node, node.args[0])
         input_qspec = (
-            SharedQuantizationSpec(input_node)
-            if is_output_annotated(input_node)
+            SharedQuantizationSpec(node.args[0])  # type: ignore[arg-type]
+            if is_output_annotated(node.args[0])  # type: ignore[arg-type]
             else input_act_qspec
         )
         quant_properties.quant_inputs = [_QuantProperty(0, input_qspec)]
         quant_properties.quant_output = _QuantProperty(
             0,
-            SharedQuantizationSpec((input_node, node)),
+            SharedQuantizationSpec((node.args[0], node)),  # type: ignore[arg-type]
         )
     elif node.target in (
         torch.ops.aten.cat.default,
@@ -554,24 +550,26 @@ def any_or_hardtanh_min_zero(n: Node):
             )
         if len(node.args[0]) == 0:
             raise ValueError("Expected non-empty list for node.args[0]")
-        inputs = [ensure_type(Node, element) for element in node.args[0]]
-        shared_qspec = SharedQuantizationSpec((inputs[0], node))
+
+        shared_qspec = SharedQuantizationSpec((node.args[0][0], node))  # type: ignore[arg-type]
         quant_properties.quant_inputs = [
             _QuantProperty(
                 0,
-                [input_act_qspec if n == inputs[0] else shared_qspec for n in inputs],
+                [
+                    input_act_qspec if n == node.args[0][0] else shared_qspec  # type: ignore[misc]
+                    for n in node.args[0]
+                ],
             )
         ]
         quant_properties.quant_output = _QuantProperty(0, shared_qspec)
     elif node.target in _one_to_one:
         quant_properties.quant_inputs = [_QuantProperty(0, input_act_qspec)]
         quant_properties.quant_output = _QuantProperty(0, output_act_qspec)
     elif node.target in _one_to_one_shared_input_qspec:
-        input_node = ensure_type(Node, node.args[0])
         quant_properties.quant_inputs = [_QuantProperty(0, input_act_qspec)]
         quant_properties.quant_output = _QuantProperty(
             0,
-            SharedQuantizationSpec((input_node, node)),
+            SharedQuantizationSpec((node.args[0], node)),  # type: ignore[arg-type]
         )
     elif node.target in [
         torch.ops.aten.eq.Tensor,
@@ -580,8 +578,7 @@ def any_or_hardtanh_min_zero(n: Node):
         torch.ops.aten.le.Tensor,
         torch.ops.aten.lt.Tensor,
     ]:
-        input_node = ensure_type(Node, node.args[0])
-        shared_qspec = SharedQuantizationSpec((input_node, node))
+        shared_qspec = SharedQuantizationSpec((node.args[0], node))  # type: ignore[arg-type]
         quant_properties.quant_inputs = [
             _QuantProperty(0, input_act_qspec),
             _QuantProperty(
@@ -599,10 +596,9 @@ def any_or_hardtanh_min_zero(n: Node):
         quant_properties.quant_inputs = []
         quant_properties.quant_output = _QuantProperty(0, output_act_qspec)
     elif node.target in [operator.getitem]:
-        input_node = ensure_type(Node, node.args[0])
-        if not is_output_annotated(input_node):
+        if not is_output_annotated(node.args[0]):  # type: ignore[arg-type]
             return None
-        shared_qspec = SharedQuantizationSpec(input_node)
+        shared_qspec = SharedQuantizationSpec(node.args[0])  # type: ignore[arg-type]
         quant_properties.quant_inputs = [_QuantProperty(0, shared_qspec)]
         quant_properties.quant_output = _QuantProperty(0, shared_qspec)
     else:

backends/arm/quantizer/quantization_annotator.py

@@ -604,9 +604,9 @@ def run_transform_for_annotation_pipeline(
         # We need to clone the artifact in order to ensure that the state_dict is preserved after passes are run.
         artifact = self.get_artifact(stage)
         if self.cur == StageType.EXPORT:
-            new_gm = ArmPassManager(
-                self.compile_spec.tosa_spec
-            ).transform_for_annotation_pipeline(graph_module=artifact.graph_module)
+            new_gm = ArmPassManager(self.compile_spec.tosa_spec).transform_for_annotation_pipeline(  # type: ignore[arg-type]
+                graph_module=artifact.graph_module
+            )
         else:
             raise RuntimeError("Can only run passes on Export stage.")
         _copy_module(artifact.graph_module, new_gm)