Arm backend: Fix mypy warnings in test/misc

Signed-off-by: [email protected]
Change-Id: I9a5dc748b3218eb7dd3ac92a7111f9404b905415

Author: perheld

backends/arm/test/misc/test_debug_feats.py

@@ -50,17 +50,19 @@ def forward(self, x):
 
 
 def _tosa_FP_pipeline(module: torch.nn.Module, test_data: input_t1, dump_file=None):
-
-    pipeline = TosaPipelineFP[input_t1](module, test_data, [], [])
+    aten_ops: list[str] = []
+    exir_ops: list[str] = []
+    pipeline = TosaPipelineFP[input_t1](module, test_data, aten_ops, exir_ops)
     pipeline.dump_artifact("to_edge_transform_and_lower")
     pipeline.dump_artifact("to_edge_transform_and_lower", suffix=dump_file)
     pipeline.pop_stage("run_method_and_compare_outputs")
     pipeline.run()
 
 
 def _tosa_INT_pipeline(module: torch.nn.Module, test_data: input_t1, dump_file=None):
-
-    pipeline = TosaPipelineINT[input_t1](module, test_data, [], [])
+    aten_ops: list[str] = []
+    exir_ops: list[str] = []
+    pipeline = TosaPipelineINT[input_t1](module, test_data, aten_ops, exir_ops)
     pipeline.dump_artifact("to_edge_transform_and_lower")
     pipeline.dump_artifact("to_edge_transform_and_lower", suffix=dump_file)
     pipeline.pop_stage("run_method_and_compare_outputs")
@@ -105,11 +107,13 @@ def test_INT_artifact(test_data: input_t1):
 
 @common.parametrize("test_data", Linear.inputs)
 def test_numerical_diff_print(test_data: input_t1):
+    aten_ops: list[str] = []
+    exir_ops: list[str] = []
     pipeline = TosaPipelineINT[input_t1](
         Linear(),
         test_data,
-        [],
-        [],
+        aten_ops,
+        exir_ops,
         custom_path="diff_print_test",
     )
     pipeline.pop_stage("run_method_and_compare_outputs")
@@ -131,7 +135,9 @@ def test_numerical_diff_print(test_data: input_t1):
 
 @common.parametrize("test_data", Linear.inputs)
 def test_dump_ops_and_dtypes(test_data: input_t1):
-    pipeline = TosaPipelineINT[input_t1](Linear(), test_data, [], [])
+    aten_ops: list[str] = []
+    exir_ops: list[str] = []
+    pipeline = TosaPipelineINT[input_t1](Linear(), test_data, aten_ops, exir_ops)
     pipeline.pop_stage("run_method_and_compare_outputs")
     pipeline.add_stage_after("quantize", pipeline.tester.dump_dtype_distribution)
     pipeline.add_stage_after("quantize", pipeline.tester.dump_operator_distribution)
@@ -149,7 +155,9 @@ def test_dump_ops_and_dtypes(test_data: input_t1):
 
 @common.parametrize("test_data", Linear.inputs)
 def test_dump_ops_and_dtypes_parseable(test_data: input_t1):
-    pipeline = TosaPipelineINT[input_t1](Linear(), test_data, [], [])
+    aten_ops: list[str] = []
+    exir_ops: list[str] = []
+    pipeline = TosaPipelineINT[input_t1](Linear(), test_data, aten_ops, exir_ops)
     pipeline.pop_stage("run_method_and_compare_outputs")
     pipeline.add_stage_after("quantize", pipeline.tester.dump_dtype_distribution, False)
     pipeline.add_stage_after(
@@ -177,7 +185,9 @@ def test_collate_tosa_INT_tests(test_data: input_t1):
     # Set the environment variable to trigger the collation of TOSA tests
     os.environ["TOSA_TESTCASES_BASE_PATH"] = "test_collate_tosa_tests"
     # Clear out the directory
-    pipeline = TosaPipelineINT[input_t1](Linear(), test_data, [], [])
+    aten_ops: list[str] = []
+    exir_ops: list[str] = []
+    pipeline = TosaPipelineINT[input_t1](Linear(), test_data, aten_ops, exir_ops)
     pipeline.pop_stage("run_method_and_compare_outputs")
     pipeline.run()
 
@@ -197,11 +207,13 @@ def test_collate_tosa_INT_tests(test_data: input_t1):
 @common.parametrize("test_data", Linear.inputs)
 def test_dump_tosa_debug_json(test_data: input_t1):
     with tempfile.TemporaryDirectory() as tmpdir:
+        aten_ops: list[str] = []
+        exir_ops: list[str] = []
         pipeline = TosaPipelineINT[input_t1](
             module=Linear(),
             test_data=test_data,
-            aten_op=[],
-            exir_op=[],
+            aten_op=aten_ops,
+            exir_op=exir_ops,
             custom_path=tmpdir,
             tosa_debug_mode=ArmCompileSpec.DebugMode.JSON,
         )
@@ -228,11 +240,13 @@ def test_dump_tosa_debug_json(test_data: input_t1):
 @common.parametrize("test_data", Linear.inputs)
 def test_dump_tosa_debug_tosa(test_data: input_t1):
     with tempfile.TemporaryDirectory() as tmpdir:
+        aten_ops: list[str] = []
+        exir_ops: list[str] = []
         pipeline = TosaPipelineINT[input_t1](
             module=Linear(),
             test_data=test_data,
-            aten_op=[],
-            exir_op=[],
+            aten_op=aten_ops,
+            exir_op=exir_ops,
             custom_path=tmpdir,
             tosa_debug_mode=ArmCompileSpec.DebugMode.TOSA,
         )
@@ -248,7 +262,9 @@ def test_dump_tosa_debug_tosa(test_data: input_t1):
 
 @common.parametrize("test_data", Linear.inputs)
 def test_dump_tosa_ops(caplog, test_data: input_t1):
-    pipeline = TosaPipelineINT[input_t1](Linear(), test_data, [], [])
+    aten_ops: list[str] = []
+    exir_ops: list[str] = []
+    pipeline = TosaPipelineINT[input_t1](Linear(), test_data, aten_ops, exir_ops)
     pipeline.pop_stage("run_method_and_compare_outputs")
     pipeline.dump_operator_distribution("to_edge_transform_and_lower")
     pipeline.run()
@@ -267,8 +283,10 @@ def forward(self, x):
 @common.parametrize("test_data", Add.inputs)
 @common.XfailIfNoCorstone300
 def test_fail_dump_tosa_ops(caplog, test_data: input_t1):
+    aten_ops: list[str] = []
+    exir_ops: list[str] = []
     pipeline = EthosU55PipelineINT[input_t1](
-        Add(), test_data, [], [], use_to_edge_transform_and_lower=True
+        Add(), test_data, aten_ops, exir_ops, use_to_edge_transform_and_lower=True
     )
     pipeline.dump_operator_distribution("to_edge_transform_and_lower")
     pipeline.run()

backends/arm/test/misc/test_debug_hook.py

@@ -5,11 +5,14 @@
 
 from dataclasses import dataclass
 from types import SimpleNamespace
+from typing import cast
 
 from executorch.backends.arm.common.arm_compile_spec import ArmCompileSpec
 from executorch.backends.arm.debug.schema import DebugHook, DebugSchema
 from executorch.backends.arm.test import common
 
+from torch.fx import Node
+
 
 @dataclass
 class DebugHookTestCase:
@@ -95,9 +98,9 @@ def create_mock_node_3():
     return fx_node_mock
 
 
-def _compare_tosa_and_schema(debug_event: DebugSchema, tosa_op):
+def _compare_tosa_and_schema(debug_event: DebugSchema, tosa_op: str) -> None:
     tosa_info = debug_event.tosa_info
-
+    assert tosa_info is not None
     assert tosa_info.node_name == tosa_op
 
     # The mapping between op_ids to operator names could change
@@ -159,7 +162,7 @@ def _compare_node_and_schema(debug_event: DebugSchema, mocked_node):
 @common.parametrize("test_data", TESTCASES)
 def test_debug_hook_add_json(test_data: DebugHookTestCase):
     hook = DebugHook(ArmCompileSpec.DebugMode.JSON)
-    hook.add(test_data.mock_node, test_data.tosa_op, test_data.op_id)
+    hook.add(cast(Node, test_data.mock_node), test_data.tosa_op, test_data.op_id)
 
     debug_events = hook._debug_events
     assert len(debug_events) == test_data.expected_events
@@ -172,7 +175,7 @@ def test_debug_hook_add_json(test_data: DebugHookTestCase):
 @common.parametrize("test_data", TESTCASES)
 def test_debug_hook_add_tosa(test_data: DebugHookTestCase):
     hook = DebugHook(ArmCompileSpec.DebugMode.TOSA)
-    hook.add(test_data.mock_node, test_data.tosa_op, test_data.op_id)
+    hook.add(cast(Node, test_data.mock_node), test_data.tosa_op, test_data.op_id)
 
     debug_events = hook._debug_events
     assert len(debug_events) == test_data.expected_events

backends/arm/test/misc/test_dim_order.py

@@ -96,28 +96,32 @@ def forward(self, x):
 
 
 @common.parametrize("module", test_modules)
-def test_dim_order_tosa_FP(module):
-    pipeline = TosaPipelineFP[input_t1](module(), module.inputs, [])
+def test_dim_order_tosa_FP(module) -> None:
+    aten_ops: list[str] = []
+    pipeline = TosaPipelineFP[input_t1](module(), module.inputs, aten_ops)
     pipeline.run()
 
 
 @common.parametrize("module", test_modules)
-def test_dim_order_tosa_INT(module):
+def test_dim_order_tosa_INT(module) -> None:
+    aten_ops: list[str] = []
     pipeline = TosaPipelineINT[input_t1](
-        module(), module.inputs, [], symmetric_io_quantization=True
+        module(), module.inputs, aten_ops, symmetric_io_quantization=True
     )
     pipeline.run()
 
 
 @common.XfailIfNoCorstone300
 @common.parametrize("module", test_modules)
-def test_dim_order_u55_INT(module):
-    pipeline = EthosU55PipelineINT[input_t1](module(), module.inputs, [])
+def test_dim_order_u55_INT(module) -> None:
+    aten_ops: list[str] = []
+    pipeline = EthosU55PipelineINT[input_t1](module(), module.inputs, aten_ops)
     pipeline.run()
 
 
 @common.XfailIfNoCorstone320
 @common.parametrize("module", test_modules)
-def test_dim_order_u85_INT(module):
-    pipeline = EthosU85PipelineINT[input_t1](module(), module.inputs, [])
+def test_dim_order_u85_INT(module) -> None:
+    aten_ops: list[str] = []
+    pipeline = EthosU85PipelineINT[input_t1](module(), module.inputs, aten_ops)
     pipeline.run()

backends/arm/test/misc/test_lifted_tensor.py

@@ -4,7 +4,8 @@
 # LICENSE file in the root directory of this source tree.
 
 import operator
-from typing import Tuple, Union
+from collections.abc import Callable
+from typing import Union
 
 import torch
 from executorch.backends.arm.test import common
@@ -15,12 +16,22 @@
 from executorch.backends.test.harness.stages import StageType
 
 
-input_t1 = Tuple[torch.Tensor]
+LiftedTensorInputs = tuple[torch.Tensor, int]
+LiftedTensorCase = tuple[
+    Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
+    LiftedTensorInputs,
+]
+LiftedScalarTensorInputs = tuple[torch.Tensor, ...]
+LiftedScalarTensorCase = tuple[
+    Callable[[torch.Tensor, Union[float, int, torch.Tensor]], torch.Tensor],
+    LiftedScalarTensorInputs,
+    Union[float, int, torch.Tensor],
+]
 
 
 class LiftedTensor(torch.nn.Module):
 
-    test_data = {
+    test_data: dict[str, LiftedTensorCase] = {
         # test_name: (operator, test_data, length)
         "add": (operator.add, (torch.randn(2, 2), 2)),
         "truediv": (operator.truediv, (torch.ones(2, 2), 2)),
@@ -39,7 +50,7 @@ def forward(self, x: torch.Tensor, length) -> torch.Tensor:
 
 
 class LiftedScalarTensor(torch.nn.Module):
-    test_data = {
+    test_data: dict[str, LiftedScalarTensorCase] = {
         # test_name: (operator, test_data)
         "add": (operator.add, (torch.randn(2, 2),), 1.0),
         "truediv": (operator.truediv, (torch.randn(4, 2),), 1.0),
@@ -60,14 +71,14 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 
 @common.parametrize("test_data", LiftedTensor.test_data)
-def test_partition_lifted_tensor_tosa_FP(test_data: input_t1):
-    op = test_data[0]
-    data = test_data[1:]
+def test_partition_lifted_tensor_tosa_FP(test_data: LiftedTensorCase) -> None:
+    op, inputs = test_data
     module = LiftedTensor(op)
-    pipeline = TosaPipelineFP[input_t1](
+    aten_ops: list[str] = []
+    pipeline = TosaPipelineFP[LiftedTensorInputs](
         module,
-        *data,
-        [],
+        inputs,
+        aten_ops,
         exir_op=[],
         use_to_edge_transform_and_lower=False,
     )
@@ -81,14 +92,14 @@ def test_partition_lifted_tensor_tosa_FP(test_data: input_t1):
 
 
 @common.parametrize("test_data", LiftedTensor.test_data)
-def test_partition_lifted_tensor_tosa_INT(test_data: input_t1):
-    op = test_data[0]
-    data = test_data[1:]
+def test_partition_lifted_tensor_tosa_INT(test_data: LiftedTensorCase) -> None:
+    op, inputs = test_data
     module = LiftedTensor(op)
-    pipeline = TosaPipelineINT[input_t1](
+    aten_ops: list[str] = []
+    pipeline = TosaPipelineINT[LiftedTensorInputs](
         module,
-        *data,
-        [],
+        inputs,
+        aten_ops,
         exir_op=[],
         use_to_edge_transform_and_lower=False,
     )
@@ -102,29 +113,33 @@ def test_partition_lifted_tensor_tosa_INT(test_data: input_t1):
 
 
 @common.parametrize("test_data", LiftedScalarTensor.test_data)
-def test_partition_lifted_scalar_tensor_tosa_FP(test_data: input_t1):
-    op = test_data[0]
-    data = test_data[1:]
-    module = LiftedScalarTensor(op, data[-1])
-    pipeline = TosaPipelineFP[input_t1](
+def test_partition_lifted_scalar_tensor_tosa_FP(
+    test_data: LiftedScalarTensorCase,
+) -> None:
+    op, tensor_inputs, scalar_arg = test_data
+    module = LiftedScalarTensor(op, scalar_arg)
+    aten_ops: list[str] = []
+    pipeline = TosaPipelineFP[LiftedScalarTensorInputs](
         module,
-        data[0],
-        [],
+        tensor_inputs,
+        aten_ops,
         exir_op=[],
         use_to_edge_transform_and_lower=False,
     )
     pipeline.run()
 
 
 @common.parametrize("test_data", LiftedScalarTensor.test_data)
-def test_partition_lifted_scalar_tensor_tosa_INT(test_data: input_t1):
-    op = test_data[0]
-    data = test_data[1:]
-    module = LiftedScalarTensor(op, data[-1])
-    pipeline = TosaPipelineINT[input_t1](
+def test_partition_lifted_scalar_tensor_tosa_INT(
+    test_data: LiftedScalarTensorCase,
+) -> None:
+    op, tensor_inputs, scalar_arg = test_data
+    module = LiftedScalarTensor(op, scalar_arg)
+    aten_ops: list[str] = []
+    pipeline = TosaPipelineINT[LiftedScalarTensorInputs](
         module,
-        data[0],
-        [],
+        tensor_inputs,
+        aten_ops,
         exir_op=[],
         use_to_edge_transform_and_lower=False,
     )

backends/arm/test/misc/test_multiple_delegates.py

@@ -29,7 +29,11 @@ def forward(self, x: torch.Tensor, y: torch.Tensor):
 
 @common.parametrize("test_data", MultipleDelegatesModule.inputs)
 def test_tosa_FP_pipeline(test_data: input_t1):
-    pipeline = TosaPipelineFP[input_t1](MultipleDelegatesModule(), test_data, [], [])
+    aten_ops: list[str] = []
+    exir_ops: list[str] = []
+    pipeline = TosaPipelineFP[input_t1](
+        MultipleDelegatesModule(), test_data, aten_ops, exir_ops
+    )
     pipeline.change_args(
         "check_count.exir", {"torch.ops.higher_order.executorch_call_delegate": 2}
     )
@@ -38,8 +42,10 @@ def test_tosa_FP_pipeline(test_data: input_t1):
 
 @common.parametrize("test_data", MultipleDelegatesModule.inputs)
 def test_tosa_INT_pipeline(test_data: input_t1):
+    aten_ops: list[str] = []
+    exir_ops: list[str] = []
     pipeline = TosaPipelineINT[input_t1](
-        MultipleDelegatesModule(), test_data, [], [], qtol=1
+        MultipleDelegatesModule(), test_data, aten_ops, exir_ops, qtol=1
     )
     pipeline.change_args(
         "check_count.exir", {"torch.ops.higher_order.executorch_call_delegate": 2}