Arm backend: Fix meandim decomp bug

Wrong indexing caused the 4th dim to be skipped when
reshaping 5D tensors. Additionally add unittest which
covers this case.

Signed-off-by: Adrian Lundell <[email protected]>
Change-Id: Ia7addeec7e4c01b2afde41e18a0b6ce932832238

Author: AdrianLundell

backends/arm/_passes/decompose_meandim_pass.py

@@ -105,7 +105,7 @@ def call_operator(self, op, args, kwargs, meta):
 
         # Reshape back to 5D if necessary
         if len(input_shape) > 4:
-            original_dims = input_shape[0:-4]
+            original_dims = input_shape[0:-3]
             temp_shape = list(x.data.shape)[1:]
             temp_shape = original_dims + temp_shape
             dims_to_reduce = [dim + len(original_dims) - 1 for dim in dims_to_reduce]

backends/arm/test/ops/test_mean_dim.py

@@ -210,6 +210,11 @@ class MeanDim(torch.nn.Module):
             (1, 2),
             False,
         ),
+        "rank5_2": lambda: (
+            torch.rand(1, 4, 7, 3, 2),
+            (2),
+            False,
+        ),
         "u55_avg_pool_not_supported": lambda: (
             torch.rand(1, 1, 1, 257),
             (0, 1, 2, 3),
@@ -255,6 +260,7 @@ def test_mean_dim_tosa_BI(test_data):
     "rank5_01234": "Rank 5 graph input currently not supported in EthosUBackend (passes since CHW are all averaged over so data order does not matter in this case)",
     "rank5_234": "Rank 5 graph input currently not supported in EthosUBackend (passes since CHW are all averaged over so data order does not matter in this case)",
     "rank5_12": "Rank 5 graph input currently not supported in EthosUBackend",
+    "rank5_2": "Rank 5 graph input currently not supported in EthosUBackend",
 }