Read reduce keepdim when it is serialized as Bool (#24379)

Summary:

mean.dim and sum.dim_IntList read keepdim only when it was serialized as Int, but it is serialized as Bool, so a dynamic resize dropped the kept dim.
keepdim is now read as Bool or Int (anything else throws), and the output rank is checked against it at build time.
This unblocks normalization built from mean(keepdim=True) in dynamic-length models.

Reviewed By: SS-JIA

Differential Revision: D127010334

Pull Request resolved: #24379

Author: psiddh

backends/webgpu/runtime/WebGPUUtils.h

@@ -301,6 +301,19 @@ inline float scalar_or(WebGPUGraph& graph, int id, float fallback) {
   throw std::runtime_error("scalar_or: expected Double, Int, or None");
 }
 
+// Reads a `bool` arg, serialized as Bool; Int is also accepted. Any other
+// type, None included, throws instead of defaulting.
+inline bool get_bool_or_int(WebGPUGraph& graph, int id, const char* what) {
+  const auto vt = graph.get_value_type(id);
+  if (vt == WebGPUGraph::ValueType::Bool) {
+    return graph.get_bool(id);
+  }
+  if (vt == WebGPUGraph::ValueType::Int) {
+    return graph.get_int(id) != 0;
+  }
+  throw std::runtime_error(std::string(what) + " must be a Bool or Int");
+}
+
 // Resolves an optional tensor arg (bias/mask/affine) to the binding to use:
 // the real tensor's buffer+size if present, else a zero-filled dummy. Caller
 // releases `.owned_dummy` (if non-null) after the dispatch is recorded.

backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp

@@ -113,18 +113,8 @@ void embedding_q4gsw_impl(WebGPUGraph& graph, const std::vector<int>& args) {
 
   // is_linear_weight selects the nibble packing (false: even dim = high nibble;
   // true: even dim = low nibble). The shader handles both via a uniform.
-  bool is_linear = false;
-  if (graph.get_value_type(is_linear_weight_id) ==
-      WebGPUGraph::ValueType::Bool) {
-    is_linear = graph.get_bool(is_linear_weight_id);
-  } else if (
-      graph.get_value_type(is_linear_weight_id) ==
-      WebGPUGraph::ValueType::Int) {
-    is_linear = graph.get_int(is_linear_weight_id) != 0;
-  } else {
-    throw std::runtime_error(
-        "WebGPU embedding_q4gsw: is_linear_weight must be Bool or Int");
-  }
+  const bool is_linear = utils::get_bool_or_int(
+      graph, is_linear_weight_id, "WebGPU embedding_q4gsw: is_linear_weight");
 
   if (weight.dims.size() < 2 || scales.dims.size() < 2 || out.dims.empty() ||
       indices.dims.empty()) {

backends/webgpu/runtime/ops/reduce/Reduce.cpp

@@ -96,10 +96,8 @@ void reduce_impl(
   const auto& in = graph.get_tensor(in_id);
   const auto& out = graph.get_tensor(out_id);
 
-  bool keepdim = false;
-  if (graph.get_value_type(keepdim_id) == WebGPUGraph::ValueType::Int) {
-    keepdim = graph.get_int(keepdim_id) != 0;
-  }
+  const bool keepdim =
+      utils::get_bool_or_int(graph, keepdim_id, "WebGPU reduce: keepdim");
 
   if (in.dims.empty()) {
     throw std::runtime_error("WebGPU reduce: scalar input unsupported");
@@ -119,6 +117,12 @@ void reduce_impl(
   if (outer == 0 || r == 0 || inner == 0) {
     throw std::runtime_error("WebGPU reduce: zero-sized reduction");
   }
+  const size_t expected_out_rank =
+      keepdim ? in.dims.size() : in.dims.size() - reduce_dims.size();
+  if (out.dims.size() != expected_out_rank) {
+    throw std::runtime_error(
+        "WebGPU reduce: output rank does not match keepdim");
+  }
 
   uint64_t in_numel = 1;
   for (int64_t d : in.dims) {

backends/webgpu/test/native/test_dynamic_shape.cpp

@@ -20,7 +20,7 @@
 //   M  dyn_sigmoid   N dyn_select (select_copy(0,-1), dynamic S)
 //   O  ONE dyn_conv1d graph reused across live input lengths
 //   P  ONE dyn_gelu graph reused above -> at -> above the old 1D dispatch cap
-//   Q  dyn_layer_norm(_rows) (reused graphs)
+//   Q  dyn_layer_norm(_rows)   R dyn_mean_keepdim (reused graphs)
 // .pte + goldens from test/ops/dynamic_shape/test_dynamic_shape_export.py.
 //
 // Artifacts dir: $WEBGPU_DYNAMIC_SHAPE_DIR, else argv[1], else
@@ -1933,6 +1933,24 @@ TEST(DynamicShape, LayerNormRowGridReusedGraph) {
   }
 }
 
+// R: mean.dim(keepdim=True) over the hidden dim ([1,1,S,1]) and the dynamic S
+// dim ([1,1,1,kHidden]), x - mean(x, -1, keepdim=True), and a keepdim=False
+// mean ([1,1,S]) on ONE loaded graph. keepdim is serialized as a Bool, so a
+// resize that ignored it would drop the kept dim.
+TEST(DynamicShape, MeanKeepdimReusedGraph) {
+  Module m(g_dir + "/dyn_mean_keepdim.pte");
+  ASSERT_EQ(m.load_forward(), Error::Ok) << "load dyn_mean_keepdim.pte";
+  for (int s : {7, 1, 125, 128, 7}) {
+    check_outputs(
+        m,
+        "dyn_mean_keepdim",
+        s,
+        {1, 1, s, kHidden},
+        {{1, 1, s, 1}, {1, 1, 1, kHidden}, {1, 1, s, kHidden}, {1, 1, s}},
+        1e-5f);
+  }
+}
+
 TEST(DynamicShape, SwiGluReusedGraph) {
   Module module(g_dir + "/dyn_swiglu.pte");
   ASSERT_EQ(module.load_forward(), Error::Ok) << "load dyn_swiglu.pte";

backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py

@@ -118,6 +118,15 @@ def forward(self, x: torch.Tensor):
         return (out + x if self.residual else out), mean, rstd
 
 
+class MeanKeepdimModule(torch.nn.Module):
+    """mean.dim with keepdim=True over the hidden dim and the dynamic S dim, a
+    broadcast consumer of the kept dim, and a keepdim=False mean."""
+
+    def forward(self, x: torch.Tensor):
+        m = x.mean(dim=-1, keepdim=True)
+        return m, x.mean(dim=2, keepdim=True), x - m, x.mean(dim=-1)
+
+
 class SwiGluModule(torch.nn.Module):
     def __init__(
         self,
@@ -374,6 +383,18 @@ def export_dynamic_layer_norm_cases(out_dir: str, s_dim) -> None:
     )
 
 
+def export_dynamic_mean_keepdim_cases(out_dir: str, s_dim) -> None:
+    _export_multi_output_case(
+        MeanKeepdimModule(),
+        "dyn_mean_keepdim",
+        out_dir,
+        lambda s: (1, 1, s, HIDDEN),
+        2,
+        s_dim,
+        (MAXS, 125, 7, 1),
+    )
+
+
 def export_dynamic_gelu_boundary_cases(out_dir: str) -> None:
     """Write a dynamic GELU fixture crossing the old 1D dispatch cap."""
     os.makedirs(out_dir, exist_ok=True)
@@ -507,6 +528,9 @@ def export_dynamic_shape_cases(out_dir: str) -> None:
     # and with a dynamic row count past the 1D grid limit.
     export_dynamic_layer_norm_cases(out_dir, s_dim)
 
+    # 2h.2) mean.dim(keepdim=True/False) over the hidden dim and the DYNAMIC S.
+    export_dynamic_mean_keepdim_cases(out_dir, s_dim)
+
     # 2i) Dynamic SwiGLU graph patterns. The Llama-width fixture forces a 2D
     # fused dispatch at M=512; compact variants isolate commutative matching and
     # graph/intermediate ownership guards around same-input q4 projections.
@@ -1713,6 +1737,11 @@ def test_export_dynamic_rms(self) -> None:
                 "dyn_layer_norm_rows.pte",
                 "dyn_layer_norm_rows.S70000.input.bin",
                 "dyn_layer_norm_rows.S1000.golden2.bin",
+                "dyn_mean_keepdim.pte",
+                "dyn_mean_keepdim.S125.input.bin",
+                "dyn_mean_keepdim.S125.golden0.bin",
+                "dyn_mean_keepdim.S125.golden3.bin",
+                "dyn_mean_keepdim.S7.golden2.bin",
                 "dyn_linear_bk64.pte",
                 "dyn_linear_bk64.S512.input.bin",
                 "dyn_linear_bk64.S512.golden.bin",