From 1fca25251805e1ea094b9592efc46955195d03f9 Mon Sep 17 00:00:00 2001
From: Julian Ng-Thow-Hing <juliannth@meta.com>
Date: Thu, 16 Jul 2026 14:49:24 -0700
Subject: [PATCH] Update

[ghstack-poisoned]
---
 backends/webgpu/test/op_tests/cases.py      | 35 +++++++++++++++++
 backends/webgpu/test/ops/test_batch_norm.py | 42 +++++++++++++++++++++
 2 files changed, 77 insertions(+)
 create mode 100644 backends/webgpu/test/ops/test_batch_norm.py

diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py
index 2bb6ff68165..efddc02be40 100644
--- a/backends/webgpu/test/op_tests/cases.py
+++ b/backends/webgpu/test/op_tests/cases.py
@@ -723,6 +723,41 @@ def _upsample_bilinear2d_suite() -> WebGPUTestSuite:
         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,
diff --git a/backends/webgpu/test/ops/test_batch_norm.py b/backends/webgpu/test/ops/test_batch_norm.py
new file mode 100644
index 00000000000..9b604f1893a
--- /dev/null
+++ b/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)