Arm backend: fix meandim when dim = None

This is a valid argument, but the pass did not support it.

Signed-off-by: Erik Lundell <[email protected]>
Change-Id: I27d0cf1732b07b9c4b0aa100e730bd1580716dc6

Author: Erik-Lundell

backends/arm/_passes/decompose_meandim_pass.py

@@ -94,6 +94,8 @@ def call_operator(self, op, args, kwargs, meta):
         input_shape = list(x.data.shape)
         output_shape = list(meta["val"].shape)
         dims_to_reduce = get_node_arg(args, 1)
+        if dims_to_reduce is None:
+            dims_to_reduce = range(len(input_shape))
         dims_to_reduce = [dim % len(input_shape) for dim in dims_to_reduce]
         dims_to_reduce = [dim for dim in dims_to_reduce if input_shape[dim] != 1]
 

backends/arm/test/ops/test_mean_dim.py

@@ -115,7 +115,7 @@ class MeanDim(torch.nn.Module):
     test_data_suite: dict[str, tuple] = {
         "rank_1_keepdim": lambda: (
             torch.rand(7),
-            (0),
+            0,
             True,
         ),
         "rank_2_keepdim": lambda: (
@@ -168,6 +168,11 @@ class MeanDim(torch.nn.Module):
             (0, 1, 2, 3),
             True,
         ),
+        "rand_none_keepdim": lambda: (
+            torch.rand(1, 5, 7, 3),
+            None,
+            True,
+        ),
         "rank_1": lambda: (
             torch.rand(7),
             (-1),
@@ -280,7 +285,6 @@ def test_mean_dim_tosa_INT(test_data):
         (test_data,),
         [],  # Might be sum, avgpool, or both
         symmetric_io_quantization=True,
-        custom_path="MEANDIM",
     )
     pipeline.run()