{% if usesF16 %}
enable f16;
{% endif %}
{{ env.wgsl.resourceDeclarations }}

const HIDDEN: u32 = {{ hiddenSize }}u;
const SPATIAL: u32 = {{ spatial }}u;
const CPG: u32 = {{ channelsPerGroup }}u;
const GROUPS: u32 = {{ numGroups }}u;
const WG: u32 = {{ workgroupSize }}u;
const SPLIT: u32 = {{ split }}u;
const EPSILON: f32 = {{ epsilon }};

@compute @workgroup_size(WG, 1, 1)
fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
  let row = wg.x;
  let part = wg.z;
  if (row >= params.rows) { return; }
  var pair = vec2<f32>(0.0);
  for (var p = 0u; p < SPLIT; p = p + 1u) { pair = pair + partials[row * SPLIT + p]; }
  let base = row * HIDDEN;
  let shift = f32(x[base]);
  let mean_d = pair.x / f32(HIDDEN);
  let variance = max(pair.y / f32(HIDDEN) - mean_d * mean_d, 0.0);
  let mean = shift + mean_d;
  let inv_std = inverseSqrt(variance + EPSILON);
  let group = row % GROUPS;
  let chunk = (HIDDEN + SPLIT - 1u) / SPLIT;
  let start = part * chunk;
  let end = min(start + chunk, HIDDEN);
  for (var d = start + lid.x; d < end; d = d + WG) {
    let channel = group * CPG + d / SPATIAL;
    let index = base + d;
    y[index] = {{ scalar }}((f32(x[index]) - mean) * inv_std * f32(scale[channel]) + f32(bias[channel]));
  }
}