Arm backend: Test avgpool non-square kernels

backends/arm/test/misc/test_tosa_dialect_avg_pool2d_adaptive.py

@@ -19,6 +19,27 @@
 from torch._subclasses.fake_tensor import FakeTensorMode
 
 
+def test_avg_pool2d_tosa_non_square_kernel_output_shape():
+    with TosaLoweringContext(
+        TosaSpecification.create_from_string("TOSA-1.0+FP")
+    ), FakeTensorMode() as mode:
+        x = mode.from_tensor(torch.randn((1, 20, 20, 8), dtype=torch.float32))
+        input_zp = mode.from_tensor(torch.zeros((1,), dtype=torch.float32))
+        output_zp = mode.from_tensor(torch.zeros((1,), dtype=torch.float32))
+
+        output = exir_ops.backend.tosa.AVG_POOL2D.default(
+            x,
+            input_zp,
+            output_zp,
+            [2, 3],
+            [2, 1],
+            [1, 1, 0, 0],
+            torch.float32,
+        )
+
+        assert tuple(output.shape) == (1, 11, 18, 8)
+
+
 def test_avg_pool2d_adaptive_tosa_INT():
     sample_inputs = [
         (

backends/arm/tosa/dialect/ops/avg_pool2d.py

@@ -136,7 +136,7 @@ def compute_avg_pool2d_output_shape(
     pad: List[IntLikeType] | List[int],
     op: str = "AVG_POOL2D",
 ) -> List[IntLikeType]:
-    """Compute the output shape for NCHW avg-pool."""
+    """Compute the output shape for NHWC avg-pool."""
 
     if x.dim() != 4:
         raise TosaValueError(f"{op} requires a 4D tensor, got {x.dim()}D", op=op)