[ExecuTorch][WebGPU] Op-tests for amax + amin

backends/webgpu/test/op_tests/cases.py

@@ -35,7 +35,6 @@
     CONFIGS as _CAT_CONFIGS,
 )
 from executorch.backends.webgpu.test.ops.test_minimum import MinimumModule
-from executorch.backends.webgpu.test.ops.test_pow import PowModule
 from executorch.backends.webgpu.test.ops.test_mul import (
     CONFIGS as _MUL_CONFIGS,
     MulModule,
@@ -44,6 +43,8 @@
     CONFIGS as _PERMUTE_CONFIGS,
     PermuteModule,
 )
+from executorch.backends.webgpu.test.ops.test_pow import PowModule
+from executorch.backends.webgpu.test.ops.test_reduce import AmaxModule, AminModule
 from executorch.backends.webgpu.test.ops.test_rms_norm import (
     _CASES,
     _linspace_weight,
@@ -224,6 +225,37 @@ def _pow_suite() -> WebGPUTestSuite:
     )
 
 
+def _reduce_suite(module_cls) -> WebGPUTestSuite:
+    # Last-dim reduction; both keepdim variants over a 2d and a 3d shape.
+    return WebGPUTestSuite(
+        module_factory=lambda keepdim: module_cls(keepdim),
+        cases=[
+            Case(name="keepdim_2d", construct={"keepdim": True}, inputs=((M1, M2),)),
+            Case(name="nodim_2d", construct={"keepdim": False}, inputs=((M1, M2),)),
+            Case(
+                name="keepdim_3d",
+                construct={"keepdim": True},
+                inputs=((S, S1, S2),),
+            ),
+            Case(
+                name="nodim_3d",
+                construct={"keepdim": False},
+                inputs=((S, S1, S2),),
+            ),
+        ],
+    )
+
+
+@register_op_test("amax")
+def _amax_suite() -> WebGPUTestSuite:
+    return _reduce_suite(AmaxModule)
+
+
+@register_op_test("amin")
+def _amin_suite() -> WebGPUTestSuite:
+    return _reduce_suite(AminModule)
+
+
 @register_op_test("view_copy")
 def _view_copy_suite() -> WebGPUTestSuite:
     return _fn_config_suite(ViewModule, _VIEW_CONFIGS)

backends/webgpu/test/ops/test_reduce.py

@@ -11,6 +11,10 @@
 outer/r/inner decomposition: `dim=-1` gives inner=1 (unit-stride reduction), a
 middle dim gives inner>1 (the non-unit-stride path); `keepdim` toggles whether the
 reduced dim survives in the output shape.
+
+`AmaxModule`/`AminModule` (below) are imported by `cases.py` for the amax/amin
+op-test suites. The WebGPU backend supports only the last-dim (per-row) reduction
+on buffer storage, so those reduce over `dim=-1`, mirroring Vulkan's per-row path.
 """
 
 from __future__ import annotations
@@ -123,5 +127,23 @@ def export_reduce_model(
     print(f"Exported {pte_path}; golden {golden_path}; input {input_path}")
 
 
+class AmaxModule(torch.nn.Module):
+    def __init__(self, keepdim: bool) -> None:
+        super().__init__()
+        self.keepdim = keepdim
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.amax(x, dim=-1, keepdim=self.keepdim)
+
+
+class AminModule(torch.nn.Module):
+    def __init__(self, keepdim: bool) -> None:
+        super().__init__()
+        self.keepdim = keepdim
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.amin(x, dim=-1, keepdim=self.keepdim)
+
+
 if __name__ == "__main__":
     unittest.main()