[ExecuTorch][WebGPU] Op-test for pow.Tensor_Scalar

backends/webgpu/test/op_tests/cases.py

@@ -1,3 +1,3 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
@@ -78,6 +78,8 @@
     ClampModule,
     HARDTANH_CONFIGS,
     HardtanhModule,
+    POW_SCALAR_CONFIGS,
+    PowScalarModule,
     UNARY_G1,
     UnaryModule,
 )
@@ -1042,3 +1044,19 @@
             for n, (lo, hi) in HARDTANH_CONFIGS.items()
         ],
     )
+
+
+@register_op_test("pow_scalar")
+def _pow_scalar_suite() -> WebGPUTestSuite:
+    # Positive base: WGSL pow(neg base)=NaN for any exponent; exponent baked.
+    return WebGPUTestSuite(
+        module_factory=lambda exponent: PowScalarModule(exponent),
+        cases=[
+            Case(
+                name=n,
+                construct={"exponent": e},
+                inputs=(InputSpec(shape=(M1, M2), gen=_unary_lin(0.1, 4.0)),),
+            )
+            for n, e in POW_SCALAR_CONFIGS.items()
+        ],
+    )

backends/webgpu/test/ops/test_unary_activations.py

@@ -62,6 +62,24 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 }
 
 
+class PowScalarModule(torch.nn.Module):
+    """aten.pow.Tensor_Scalar with a baked exponent (exponent → the min slot)."""
+
+    def __init__(self, exponent) -> None:
+        super().__init__()
+        self.exponent = exponent
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.pow(x, self.exponent)
+
+
+# name -> exponent construct kwarg; the suite uses a positive base to avoid NaN.
+POW_SCALAR_CONFIGS = {
+    "square": 2.0,
+    "sqrt": 0.5,
+}
+
+
 def _lin(lo: float, hi: float):
     """Deterministic linspace input of the requested shape over [lo, hi]."""