Arm backend: Fix bug in decompose linear vector norm (pytorch#11755)

The introduction of decomposition for linalg vector norm revealed a bug
that when dim is None, then all dimensions should be reduced.

Signed-off-by: Elena Zhelezina <[email protected]>

Author: wwwind

backends/arm/_passes/decompose_linalg_vector_norm_pass.py

@@ -51,10 +51,12 @@ def call_operator(self, op, args, kwargs, meta):
                 f"is not supported for linalg_vector_norm operator"
             )
 
+        # Sum over all dimensions if dim is None
         if norm_dim is None:
-            raise ValueError("The norm_dim for linalg_vector_norm is None.")
-
-        dims = [norm_dim] if isinstance(norm_dim, int) else list(norm_dim)
+            rank = input_tensor.data.dim()
+            dims = list(range(rank))
+        else:
+            dims = [norm_dim] if isinstance(norm_dim, int) else list(norm_dim)
 
         # Decomposition based on norm order.
         if norm_order == 1:

backends/arm/_passes/decompose_sum_pass.py

@@ -63,6 +63,11 @@ def call_operator(self, op, args, kwargs, meta):
             case _:
                 raise ValueError(f"Invalid number of arguments ({len(args)}) provided.")
 
+        # If dims is None, sum over all dimensions
+        if dims is None:
+            shape = input_node.data.size()
+            dims = list(range(len(shape)))
+
         view_op, sum_op = _get_sum_decomp(op)
 
         for dim in dims:

backends/arm/operators/op_sum.py

@@ -106,8 +106,6 @@ def define_node(
         if inputs[0].dtype == ts.DType.INT8:
             return super().define_node(node, tosa_graph, inputs, output)
 
-        validate_num_inputs(self.target, inputs, 3)
-
         tensor = inputs[0]
         input_shape = list(tensor.shape)
         dim = int(inputs[1].number % len(input_shape))

backends/arm/test/models/test_torch_functions.py

@@ -130,7 +130,6 @@ def test_torch_fns_MI(test_data):
         "topk": "NotImplementedError: No registered serialization name for <class 'torch.return_types.topk'> found",
         "sort": "NotImplementedError: No registered serialization name for <class 'torch.return_types.sort'> found",
         "t": "MLETORCH-855: Issue with Quantization folding.",
-        "norm": "An error occurred when running the 'KeepDimsFalseToSqueezePass' pass after the following passes:",
     },
     strict=False,
 )

backends/arm/test/ops/test_sum.py

@@ -33,6 +33,7 @@ class Sum(torch.nn.Module):
         "4d_dims_no_keep": lambda: (torch.rand(1, 1, 5, 8), 1, False),
         "4d_dim_3_keep": lambda: (torch.rand(1, 2, 3, 4), 3, True),
         "4d_dims_keep": lambda: (torch.rand(1, 2, 8, 8), [2, 3, 0], True),
+        "dim_None": lambda: (torch.rand(10), None, True),
     }
 
     def forward(self, x: torch.Tensor, dim: int, keepdim: bool):