[ExecuTorch][WebGPU] batch_norm op tests

backends/webgpu/test/op_tests/cases.py

@@ -1,3 +1,3 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
@@ -723,6 +723,41 @@
         atol=1e-4,
         rtol=1e-3,
     )
+from executorch.backends.webgpu.test.ops.test_batch_norm import (
+    _det_input as _bn_det_input,
+    BatchNorm2dModule,
+)
+
+
+@register_op_test("batch_norm")
+def _batch_norm_suite() -> WebGPUTestSuite:
+    # MODNet decoder / CNN-backbone inference batch norm (eval -> the no-training
+    # variant). Covers affine + non-affine (optional weight/bias) and an odd H*W.
+    # Only the `out` ValueList entry is compared (out_index 0).
+    return WebGPUTestSuite(
+        module_factory=lambda num_features, affine: BatchNorm2dModule(
+            num_features, affine
+        ).eval(),
+        cases=[
+            Case(
+                name="affine_c8",
+                construct={"num_features": 8, "affine": True},
+                inputs=(InputSpec(shape=(1, 8, 12, 12), gen=_bn_det_input),),
+            ),
+            Case(
+                name="no_affine_c8",
+                construct={"num_features": 8, "affine": False},
+                inputs=(InputSpec(shape=(1, 8, 12, 12), gen=_bn_det_input),),
+            ),
+            Case(
+                name="c16_odd",
+                construct={"num_features": 16, "affine": True},
+                inputs=(InputSpec(shape=(1, 16, 5, 7), gen=_bn_det_input),),
+            ),
+        ],
+        atol=1e-3,
+        rtol=1e-3,
+    )
 from executorch.backends.webgpu.test.ops.test_max_pool2d import (
     _det_input as _maxpool_det_input,
     MaxPool2dModule,

backends/webgpu/test/ops/test_batch_norm.py

@@ -0,0 +1,42 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""`aten._native_batch_norm_legit_no_training.default` module for the WebGPU op-test framework.
+
+Inference batch norm is on MODNet's decoder and many CNN backbones; `nn.BatchNorm2d`
+in eval mode lowers to `_native_batch_norm_legit_no_training`, applying the per-channel
+running-stat affine. Deterministic running stats + affine make the golden non-trivial.
+"""
+
+import torch
+
+
+class BatchNorm2dModule(torch.nn.Module):
+    def __init__(
+        self, num_features: int, affine: bool = True, eps: float = 1e-5
+    ) -> None:
+        super().__init__()
+        bn = torch.nn.BatchNorm2d(num_features, eps=eps, affine=affine)
+        with torch.no_grad():
+            bn.running_mean.copy_(torch.linspace(-1.0, 1.0, num_features))
+            bn.running_var.copy_(torch.linspace(0.5, 2.0, num_features))
+            if affine:
+                bn.weight.copy_(torch.linspace(0.5, 1.5, num_features))
+                bn.bias.copy_(torch.linspace(-0.5, 0.5, num_features))
+        bn.eval()
+        self.bn = bn
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.bn(x)
+
+
+def _det_input(shape) -> torch.Tensor:
+    # ((i % 23) - 11) / 16: exact in fp32, spans negatives through positives.
+    n = 1
+    for s in shape:
+        n *= s
+    idx = torch.arange(n, dtype=torch.int64)
+    return (((idx % 23) - 11).to(torch.float32) / 16.0).reshape(shape)