support qnn mean (dim=None)

Summary: Address mean op lower failure. When dim is not specified, it will take mean across all axes. For QNN, we need to get axes based on input shape

Differential Revision: D83520776

Author: cccclai

backends/qualcomm/builders/op_mean_dim.py

@@ -40,7 +40,16 @@ def define_node(
         )
 
         # mean dims and keep dims
-        mean_dims = cast(List[int], node.args[1])
+        rank = len(input_node.meta["val"].shape)
+        dim_arg = node.args[1]
+
+        if dim_arg is None:
+            mean_dims = list(range(rank))   # reduce over all dims
+        elif isinstance(dim_arg, int):
+            mean_dims = [dim_arg]
+        else:
+            mean_dims = list(dim_arg)
+
         mean_dims = [
             mean_dim % len(input_node.meta["val"].shape) for mean_dim in mean_dims
         ]

backends/qualcomm/tests/models.py

@@ -5,7 +5,7 @@
 # LICENSE file in the root directory of this source tree.
 
 import torch
-
+from typing import Optional, Union, Tuple, List
 
 # module with related operator only
 
@@ -1348,6 +1348,22 @@ def forward(self, x):
         return torch.mean(x, (-1, -2))
 
 
+class Mean(torch.nn.Module):
+    def __init__(
+        self,
+        dim: Optional[Union[int, Tuple[int, ...], List[int]]] = None,
+        keepdim: bool = False,
+        dtype: Optional[torch.dtype] = None,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.keepdim = keepdim
+        self.dtype = dtype
+
+    def forward(self, x):
+        return torch.mean(x, dim=self.dim, keepdim=self.keepdim, dtype=self.dtype)
+
+
 class MaskedFill(torch.nn.Module):
     def __init__(self):
         super().__init__()

backends/qualcomm/tests/test_qnn_delegate.py

@@ -1025,6 +1025,13 @@ def test_qnn_backend_mean_dim(self):
             with self.subTest(i=i):
                 self.lower_module_and_test_output(module, sample_input)
 
+    def test_qnn_backend_mean(self):
+        modules = [Mean(), Mean()]  # noqa: F405
+        sample_inputs = [(torch.randn(10, 10),), (torch.tensor([5.0]),)]
+        for i, module in enumerate(modules):
+            with self.subTest(i=i):
+                self.lower_module_and_test_output(module, sample_inputs[i])
+
     @unittest.skip("failed to lower in QNN 2.26")
     def test_qnn_backend_mha(self):
         module = MultiheadAttention()  # noqa: F405
@@ -2674,6 +2681,14 @@ def test_qnn_backend_mean_dim(self):
                 module = self.get_qdq_module(module, sample_input)
                 self.lower_module_and_test_output(module, sample_input)
 
+    def test_qnn_backend_mean(self):
+        modules = [Mean(), Mean()]  # noqa: F405
+        sample_inputs = [(torch.randn(10, 10),), (torch.tensor([5.0]),)]
+        for i, module in enumerate(modules):
+            with self.subTest(i=i):
+                module = self.get_qdq_module(module, sample_inputs[i])
+                self.lower_module_and_test_output(module, sample_inputs[i])
+
     def test_qnn_backend_mha(self):
         module = MultiheadAttention()  # noqa: F405
         sample_input = (torch.randn(1, 197, 96),)