{% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
  // per-axis workgroup fold width.
  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
  if ({{ name }} >= {{ bound }}) { return; }{% endmacro %}
{{ env.wgsl.resourceDeclarations }}
// Separate optional outputs keep the reduction within eight storage bindings.
@compute @workgroup_size({{ workgroupSize }})
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
{{ flat_index_2d(workgroupSize, "row", "params.rows", guardInline=true) }}
  let stats = row_stats[row];
  mean[row] = stats.x;
  inv_std_var[row] = stats.y;
}