{% 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 }}
@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
{{ flat_index_2d(tunables.WORKGROUP_SIZE, "index") }}
  let spatial = params.height * params.width;
  let channel = (index / spatial) % params.channels;
{% if usesF16 %}
  // f16 tensors use an f32 epsilon, so widen the normalization arithmetic to
  // f32 and narrow only on store. WGSL does not allow the mixed f16 + f32 add.
  let normalized = (f32(x[index]) - f32(input_mean[channel])) * inverseSqrt(f32(input_var[channel]) + params.epsilon);
  y[index] = {{ scalar }}(normalized * f32(scale[channel]) + f32(bias[channel]));
{% else %}
  let normalized = (x[index] - input_mean[channel]) * inverseSqrt(input_var[channel] + params.epsilon);
  y[index] = normalized * scale[channel] + bias[channel];
{% endif %}
}