onnxruntime-ep-mlx 0.2.3

//! Normalization op handlers. Faithful port of the C++ `ops/norm.cc` + `ops/norm_ext.cc`:
//!
//!   * RMSNormalization (ai.onnx opset 23+)             — mlx_fast_rms_norm
//!   * LayerNormalization (ai.onnx opset 17+)           — mlx_fast_layer_norm
//!   * SimplifiedLayerNormalization (com.microsoft)     — mlx_fast_rms_norm
//!   * SkipLayerNormalization (com.microsoft)           — residual add + mlx_fast_layer_norm
//!   * SkipSimplifiedLayerNormalization (com.microsoft) — residual add + mlx_fast_rms_norm
//!   * GroupNormalization (ai.onnx opset 21 form)       — composed mean/var/rsqrt
//!   * LpNormalization (ai.onnx)                        — composed abs/sum or square/sum/sqrt
//!   * BatchNormalization (ai.onnx, inference form)     — composed per-channel affine
//!   * LRN (ai.onnx, across-channel)                    — composed square/window-sum/power/divide
//!
//! Every handler honors the resolved input dtype (fp32/fp16/bf16) with no per-dtype branching:
//! the MLX fast norms run in whatever float dtype the input carries, and the composed paths keep a
//! matching-dtype epsilon scalar so no unwanted upcast occurs.

use crate::engine::{MlxError, NodeDesc, Src, TranslationContext};
use crate::registry::{is_mlx_float, NodeView, OpRegistration, OpRegistry, K_ANY_OPSET};
use crate::sys::mlx;
use crate::sys::ort;
use std::os::raw::c_char;

// ---- small local MLX helpers -------------------------------------------------------------------

/// A null-`ctx` `mlx_array` — the mlx-c "empty/absent" sentinel (e.g. an omitted layer-norm bias).
#[inline]
fn empty_array() -> mlx::mlx_array {
    mlx::mlx_array_ {
        ctx: std::ptr::null_mut(),
    }
}

fn mul(ctx: &mut TranslationContext, a: mlx::mlx_array, b: mlx::mlx_array) -> Result<mlx::mlx_array, MlxError> {
    ctx.binary(mlx::mlx_multiply, a, b)
}
fn add(ctx: &mut TranslationContext, a: mlx::mlx_array, b: mlx::mlx_array) -> Result<mlx::mlx_array, MlxError> {
    ctx.binary(mlx::mlx_add, a, b)
}
fn sub(ctx: &mut TranslationContext, a: mlx::mlx_array, b: mlx::mlx_array) -> Result<mlx::mlx_array, MlxError> {
    ctx.binary(mlx::mlx_subtract, a, b)
}
fn divide(ctx: &mut TranslationContext, a: mlx::mlx_array, b: mlx::mlx_array) -> Result<mlx::mlx_array, MlxError> {
    ctx.binary(mlx::mlx_divide, a, b)
}
fn rsqrt(ctx: &mut TranslationContext, a: mlx::mlx_array) -> Result<mlx::mlx_array, MlxError> {
    ctx.emit(|res, s| unsafe { mlx::mlx_rsqrt(res, a, s) })
}
fn sqrt(ctx: &mut TranslationContext, a: mlx::mlx_array) -> Result<mlx::mlx_array, MlxError> {
    ctx.emit(|res, s| unsafe { mlx::mlx_sqrt(res, a, s) })
}
fn abs(ctx: &mut TranslationContext, a: mlx::mlx_array) -> Result<mlx::mlx_array, MlxError> {
    ctx.emit(|res, s| unsafe { mlx::mlx_abs(res, a, s) })
}
fn sum_axis(ctx: &mut TranslationContext, a: mlx::mlx_array, axis: i32, keepdims: bool) -> Result<mlx::mlx_array, MlxError> {
    ctx.emit(|res, s| unsafe { mlx::mlx_sum_axis(res, a, axis, keepdims, s) })
}
fn mean_axis(ctx: &mut TranslationContext, a: mlx::mlx_array, axis: i32, keepdims: bool) -> Result<mlx::mlx_array, MlxError> {
    ctx.emit(|res, s| unsafe { mlx::mlx_mean_axis(res, a, axis, keepdims, s) })
}
fn var_axis(ctx: &mut TranslationContext, a: mlx::mlx_array, axis: i32, keepdims: bool) -> Result<mlx::mlx_array, MlxError> {
    ctx.emit(|res, s| unsafe { mlx::mlx_var_axis(res, a, axis, keepdims, 0, s) })
}

/// A 0-d scalar of dtype `dt` holding `v` (eps constant), matching the compute dtype so no unwanted
/// upcast occurs.
fn scalar_like(ctx: &mut TranslationContext, v: f32, dt: mlx::mlx_dtype) -> Result<mlx::mlx_array, MlxError> {
    let s = ctx.scalar_f32(v);
    if dt == mlx::mlx_dtype__MLX_FLOAT32 {
        Ok(s)
    } else {
        ctx.astype(s, dt)
    }
}

/// Reshape a per-channel vector `[C]` to `[1, C, 1, ..., 1]` so it broadcasts over N and spatial.
fn channel_broadcast(ctx: &mut TranslationContext, v: mlx::mlx_array, rank: usize, channels: i32) -> Result<mlx::mlx_array, MlxError> {
    let mut shape = vec![1i32; rank];
    if rank >= 2 {
        shape[1] = channels;
    }
    ctx.reshape(v, &shape)
}

fn epsilon(n: &NodeDesc, default: f32) -> f32 {
    n.floats.get("epsilon").copied().unwrap_or(default)
}

fn present(n: &NodeDesc, i: usize) -> bool {
    i < n.inputs.len() && n.inputs[i].source != Src::Absent
}

// ---- handlers ----------------------------------------------------------------------------------

/// RMSNormalization (ai.onnx opset 23+): out = rms_norm(x) * scale over the last axis.
fn rms_norm_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let x = ctx.resolve(&n.inputs[0])?;
    let g = ctx.resolve(&n.inputs[1])?;
    let eps = epsilon(n, 1e-6);
    let r = ctx.emit(|res, s| unsafe { mlx::mlx_fast_rms_norm(res, x, g, eps, s) })?;
    ctx.mark_fast("mlx_fast_rms_norm");
    ctx.bind(&n.outputs[0], r);
    Ok(())
}

/// SimplifiedLayerNormalization (com.microsoft): Y = rms_norm(X) * scale over the last axis.
fn simplified_layer_norm_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let x = ctx.resolve(&n.inputs[0])?;
    let scale = ctx.resolve(&n.inputs[1])?;
    let eps = epsilon(n, 1e-5);
    let r = ctx.emit(|res, s| unsafe { mlx::mlx_fast_rms_norm(res, x, scale, eps, s) })?;
    ctx.mark_fast("mlx_fast_rms_norm");
    ctx.bind(&n.outputs[0], r);
    Ok(())
}

/// LayerNormalization (ai.onnx opset 17+, last-axis form): Y = layer_norm(X, scale, bias, eps).
/// Only the single-output (Y) form is claimed; Mean/InvStdDev extra outputs are left to CPU.
fn layer_norm_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let x = ctx.resolve(&n.inputs[0])?;
    let scale = ctx.resolve(&n.inputs[1])?;
    let bias = if present(n, 2) {
        ctx.resolve(&n.inputs[2])?
    } else {
        empty_array()
    };
    let eps = epsilon(n, 1e-5);
    let r = ctx.emit(|res, s| unsafe { mlx::mlx_fast_layer_norm(res, x, scale, bias, eps, s) })?;
    ctx.mark_fast("mlx_fast_layer_norm");
    ctx.bind(&n.outputs[0], r);
    Ok(())
}

/// SkipLayerNormalization (com.microsoft): residual = input + skip (+ bias);
/// Y = layer_norm(residual, gamma, beta, eps). out[0]=Y; optional out[3]=residual sum.
fn skip_layer_norm_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let input = ctx.resolve(&n.inputs[0])?;
    let skip = ctx.resolve(&n.inputs[1])?;
    let gamma = ctx.resolve(&n.inputs[2])?;
    let beta = if present(n, 3) {
        ctx.resolve(&n.inputs[3])?
    } else {
        empty_array()
    };
    let mut residual = add(ctx, input, skip)?;
    if present(n, 4) {
        let bias = ctx.resolve(&n.inputs[4])?;
        residual = add(ctx, residual, bias)?;
    }
    let eps = epsilon(n, 1e-5);
    let r = ctx.emit(|res, s| unsafe { mlx::mlx_fast_layer_norm(res, residual, gamma, beta, eps, s) })?;
    ctx.mark_fast("mlx_fast_layer_norm");
    ctx.bind(&n.outputs[0], r);
    if n.outputs.len() >= 4 && !n.outputs[3].name.is_empty() {
        ctx.bind(&n.outputs[3], residual);
    }
    Ok(())
}

/// SkipSimplifiedLayerNormalization (com.microsoft): residual = input + skip;
/// out = rms_norm(residual) * gamma. out[0]=normalized, out[last]=residual.
fn skip_rms_norm_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let input = ctx.resolve(&n.inputs[0])?;
    let skip = ctx.resolve(&n.inputs[1])?;
    let gamma = ctx.resolve(&n.inputs[2])?;
    let eps = epsilon(n, 1e-6);
    let residual = add(ctx, input, skip)?;
    let norm = ctx.emit(|res, s| unsafe { mlx::mlx_fast_rms_norm(res, residual, gamma, eps, s) })?;
    ctx.mark_fast("mlx_fast_rms_norm");
    ctx.bind(&n.outputs[0], norm);
    if n.outputs.len() >= 2 {
        let last = n.outputs.len() - 1;
        ctx.bind(&n.outputs[last], residual);
    }
    Ok(())
}

/// GroupNormalization (ai.onnx opset 21 form): normalize within each of `num_groups` channel groups,
/// then apply per-channel scale/bias. X=[N,C,*S], scale/bias=[C].
fn group_norm_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let x = ctx.resolve(&n.inputs[0])?;
    let scale = ctx.resolve(&n.inputs[1])?;
    let bias = ctx.resolve(&n.inputs[2])?;
    let shape = ctx.shape_of(x);
    let rank = shape.len();
    let n_dim = shape[0];
    let c = shape[1];
    let groups = n.ints.get("num_groups").copied().unwrap_or(0) as i32;
    if groups <= 0 {
        return Err("MLX GroupNormalization requires num_groups > 0".to_string());
    }
    let eps = epsilon(n, 1e-5);

    let mut per_group: i32 = 1;
    for &d in &shape[1..rank] {
        per_group *= d;
    }
    per_group /= groups; // (C/groups) * prod(spatial)

    let grp = ctx.reshape(x, &[n_dim, groups, per_group])?;
    let mean = mean_axis(ctx, grp, 2, true)?;
    let var = var_axis(ctx, grp, 2, true)?;
    let eps_s = scalar_like(ctx, eps, ctx.dtype_of(x))?;
    let var_eps = add(ctx, var, eps_s)?;
    let inv = rsqrt(ctx, var_eps)?;
    let centered = sub(ctx, grp, mean)?;
    let normed = mul(ctx, centered, inv)?;
    let normed = ctx.reshape(normed, &shape)?;

    let sb = channel_broadcast(ctx, scale, rank, c)?;
    let bb = channel_broadcast(ctx, bias, rank, c)?;
    let scaled = mul(ctx, normed, sb)?;
    let out = add(ctx, scaled, bb)?;
    ctx.mark_composed("GroupNormalization composed (mean/var/rsqrt) — no fused last-axis norm kernel");
    ctx.bind(&n.outputs[0], out);
    Ok(())
}

/// LpNormalization (ai.onnx): Y = X / ||X||_p along `axis` (p in {1,2}, default 2; axis default -1).
fn lp_norm_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let x = ctx.resolve(&n.inputs[0])?;
    let rank = ctx.ndim(x) as i64;
    let mut axis = n.ints.get("axis").copied().unwrap_or(-1);
    if axis < 0 {
        axis += rank;
    }
    let p = n.ints.get("p").copied().unwrap_or(2);
    let axis = axis as i32;

    let norm = if p == 1 {
        let a = abs(ctx, x)?;
        sum_axis(ctx, a, axis, true)?
    } else {
        let sq = mul(ctx, x, x)?;
        let s = sum_axis(ctx, sq, axis, true)?;
        sqrt(ctx, s)?
    };
    let quot = divide(ctx, x, norm)?;
    // ONNX LpNormalization: where the norm is 0, emit 0 rather than NaN (0/0). Matches the ONNX
    // reference `np.where(norm == 0, 0, x / norm)` and ORT's CPU kernel.
    let zero = scalar_like(ctx, 0.0, ctx.dtype_of(x))?;
    let is_zero = ctx.binary(mlx::mlx_equal, norm, zero)?;
    let out = ctx.where_(is_zero, zero, quot)?;
    ctx.mark_composed("LpNormalization composed (abs/sum/sqrt/divide) — no fused norm kernel");
    ctx.bind(&n.outputs[0], out);
    Ok(())
}

/// BatchNormalization (ai.onnx, inference/spatial form): Y = (X - mean)/sqrt(var+eps) * scale + B,
/// per channel. Only the single-output inference form is claimed.
fn batch_norm_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let x = ctx.resolve(&n.inputs[0])?;
    let scale = ctx.resolve(&n.inputs[1])?;
    let b = ctx.resolve(&n.inputs[2])?;
    let mean = ctx.resolve(&n.inputs[3])?;
    let var = ctx.resolve(&n.inputs[4])?;
    let shape = ctx.shape_of(x);
    let rank = shape.len();
    let c = if rank >= 2 { shape[1] } else { shape[0] };
    let eps = epsilon(n, 1e-5);

    let eps_s = scalar_like(ctx, eps, ctx.dtype_of(x))?;
    let var_eps = add(ctx, var, eps_s)?;
    let inv = rsqrt(ctx, var_eps)?; // [C]
    let a = mul(ctx, scale, inv)?; // [C]
    let mean_a = mul(ctx, mean, a)?;
    let shift = sub(ctx, b, mean_a)?; // [C]
    let ab = channel_broadcast(ctx, a, rank, c)?;
    let shiftb = channel_broadcast(ctx, shift, rank, c)?;
    let scaled = mul(ctx, x, ab)?;
    let out = add(ctx, scaled, shiftb)?;
    ctx.mark_composed("BatchNormalization composed (rsqrt/affine) — no fused batch-norm kernel");
    ctx.bind(&n.outputs[0], out);
    Ok(())
}

// ---- LRN helpers -------------------------------------------------------------------------------

/// Pad `a` along the channel axis (axis 1) with `low`/`high` copies of `value`.
fn pad_channel(
    ctx: &mut TranslationContext,
    a: mlx::mlx_array,
    low: i32,
    high: i32,
    value: mlx::mlx_array,
) -> Result<mlx::mlx_array, MlxError> {
    if low == 0 && high == 0 {
        return Ok(a);
    }
    let axes = [1i32];
    let lo = [low];
    let hi = [high];
    let mode = b"constant\0";
    let out = ctx.emit(|res, s| unsafe {
        mlx::mlx_pad(
            res,
            a,
            axes.as_ptr(),
            1,
            lo.as_ptr(),
            1,
            hi.as_ptr(),
            1,
            value,
            mode.as_ptr() as *const c_char,
            s,
        )
    })?;
    ctx.contiguous(out)
}

/// Slice `a` along the channel axis (axis 1) to `[lo, hi)`, keeping all other axes intact.
fn slice_channel(ctx: &mut TranslationContext, a: mlx::mlx_array, lo: i32, hi: i32) -> Result<mlx::mlx_array, MlxError> {
    let shape = ctx.shape_of(a);
    let rank = shape.len();
    let mut start = vec![0i32; rank];
    let mut stop = shape;
    let stride = vec![1i32; rank];
    start[1] = lo;
    stop[1] = hi;
    ctx.emit(|res, s| unsafe {
        mlx::mlx_slice(
            res,
            a,
            start.as_ptr(),
            rank,
            stop.as_ptr(),
            rank,
            stride.as_ptr(),
            rank,
            s,
        )
    })
}

/// LRN (ai.onnx, across-channel): for input X=[N,C,*S],
///   square_sum[n,c,*] = sum over the `size`-wide channel window centered at c (clamped to [0,C-1]),
///   Y[n,c,*] = X[n,c,*] / (bias + (alpha/size) * square_sum[n,c,*])^beta.
/// The window sum is computed by zero-padding X^2 along the channel axis (so out-of-range channels
/// contribute 0, matching ONNX's clamped window) then summing `size` shifted channel slices.
fn lrn_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let x = ctx.resolve(&n.inputs[0])?;
    let shape = ctx.shape_of(x);
    let c = shape[1];
    let size = n.ints.get("size").copied().unwrap_or(1).max(1) as i32;
    let alpha = n.floats.get("alpha").copied().unwrap_or(1e-4);
    let beta = n.floats.get("beta").copied().unwrap_or(0.75);
    let bias = n.floats.get("bias").copied().unwrap_or(1.0);
    let dt = ctx.dtype_of(x);

    let x2 = mul(ctx, x, x)?;
    // Window [c - floor((size-1)/2), c + ceil((size-1)/2)] clamped to [0, C-1].
    let pad_before = (size - 1) / 2;
    let pad_after = size - 1 - pad_before;
    let zero = scalar_like(ctx, 0.0, dt)?;
    let xp = pad_channel(ctx, x2, pad_before, pad_after, zero)?; // channel length C + size - 1

    // square_sum[:, c] = sum_{k=0}^{size-1} xp[:, c + k]
    let mut square_sum = slice_channel(ctx, xp, 0, c)?;
    for k in 1..size {
        let s = slice_channel(ctx, xp, k, k + c)?;
        square_sum = add(ctx, square_sum, s)?;
    }

    let scale = scalar_like(ctx, alpha / size as f32, dt)?;
    let bias_s = scalar_like(ctx, bias, dt)?;
    let beta_s = scalar_like(ctx, beta, dt)?;
    let scaled = mul(ctx, square_sum, scale)?;
    let base = add(ctx, scaled, bias_s)?;
    let denom = ctx.binary(mlx::mlx_power, base, beta_s)?;
    let out = divide(ctx, x, denom)?;
    ctx.mark_composed("LRN composed (square/window-sum/power/divide) — no fused LRN kernel");
    ctx.bind(&n.outputs[0], out);
    Ok(())
}

// ---- claim predicates --------------------------------------------------------------------------

fn tensor_dtype(node: &NodeView, i: usize) -> Option<ort::ONNXTensorElementDataType> {
    node.input_info(i).map(|s| s.dtype)
}

/// RMSNormalization (ai.onnx): X, scale, axis == -1. fp32/fp16/bf16 (mlx_fast_rms_norm is generic).
fn rms_norm_claim(node: &NodeView) -> bool {
    if node.num_inputs() != 2 || node.num_outputs() == 0 {
        return false;
    }
    let (x, g, out) = match (node.input_info(0), tensor_dtype(node, 1), node.output_info(0)) {
        (Some(x), Some(g), Some(o)) => (x, g, o),
        _ => return false,
    };
    if !is_mlx_float(x.dtype) || g != x.dtype || out.dtype != x.dtype {
        return false;
    }
    node.int_attr("axis", -1) == -1
}

/// SimplifiedLayerNormalization (com.microsoft): X + scale, last-axis, single output.
fn simplified_layer_norm_claim(node: &NodeView) -> bool {
    if node.num_inputs() != 2 || node.num_outputs() != 1 {
        return false;
    }
    let (x, g, out) = match (node.input_info(0), tensor_dtype(node, 1), node.output_info(0)) {
        (Some(x), Some(g), Some(o)) => (x, g, o),
        _ => return false,
    };
    if !is_mlx_float(x.dtype) || g != x.dtype || out.dtype != x.dtype || x.shape.is_empty() {
        return false;
    }
    let axis = node.int_attr("axis", -1);
    axis == -1 || axis == x.shape.len() as i64 - 1
}

/// LayerNormalization: fp32/fp16/bf16 X + scale (+ optional bias), last-axis, single output (Y).
fn layer_norm_claim(node: &NodeView) -> bool {
    let nin = node.num_inputs();
    if nin < 2 || nin > 3 || node.num_outputs() != 1 {
        return false;
    }
    let (x, scale, out) = match (node.input_info(0), tensor_dtype(node, 1), node.output_info(0)) {
        (Some(x), Some(scale), Some(o)) => (x, scale, o),
        _ => return false,
    };
    if !is_mlx_float(x.dtype) || scale != x.dtype || out.dtype != x.dtype {
        return false;
    }
    if nin == 3 && node.input_present(2) {
        match tensor_dtype(node, 2) {
            Some(bias) if bias == x.dtype => {}
            _ => return false,
        }
    }
    let rank = x.shape.len() as i64;
    let axis = node.int_attr("axis", -1);
    rank > 0 && (axis == -1 || axis == rank - 1)
}

/// SkipSimplifiedLayerNormalization (com.microsoft): input, skip, gamma. fp32/fp16/bf16. 3-input.
fn skip_rms_norm_claim(node: &NodeView) -> bool {
    if node.num_inputs() != 3 || node.num_outputs() == 0 {
        return false;
    }
    let (x, out) = match (node.input_info(0), node.output_info(0)) {
        (Some(x), Some(o)) => (x, o),
        _ => return false,
    };
    if !is_mlx_float(x.dtype) || out.dtype != x.dtype {
        return false;
    }
    // The handler produces only out[0] (normalized) and the optional out[last] (residual sum);
    // reject if mean (out[1]) or inv-std (out[2]) are requested — mlx_fast_rms_norm doesn't compute
    // them, so claiming would leave those outputs unbound (mirrors skip_layer_norm_claim).
    if node.output_present(1) || node.output_present(2) {
        return false;
    }
    matches!(tensor_dtype(node, 1), Some(t) if t == x.dtype)
        && matches!(tensor_dtype(node, 2), Some(t) if t == x.dtype)
}

/// SkipLayerNormalization: input, skip, gamma (+ optional beta, bias), all same float dtype.
/// Only out[0] (Y) and optional out[3] (residual sum) are produced; mean/inv-std outputs → CPU.
fn skip_layer_norm_claim(node: &NodeView) -> bool {
    let nin = node.num_inputs();
    if nin < 3 || nin > 5 || node.num_outputs() == 0 {
        return false;
    }
    let (x, out) = match (node.input_info(0), node.output_info(0)) {
        (Some(x), Some(o)) => (x, o),
        _ => return false,
    };
    if !is_mlx_float(x.dtype) || out.dtype != x.dtype {
        return false;
    }
    for i in 1..nin {
        if !node.input_present(i) {
            continue;
        }
        match tensor_dtype(node, i) {
            Some(t) if t == x.dtype => {}
            _ => return false,
        }
    }
    // Reject if mean (out[1]) or inv-std (out[2]) are requested — we do not compute them.
    if node.output_present(1) || node.output_present(2) {
        return false;
    }
    true
}

/// GroupNormalization: X=[N,C,*S] float, scale/bias=[C], static C divisible by num_groups.
fn group_norm_claim(node: &NodeView) -> bool {
    if node.num_inputs() != 3 || node.num_outputs() != 1 {
        return false;
    }
    let (x, scale, out) = match (node.input_info(0), node.input_info(1), node.output_info(0)) {
        (Some(x), Some(scale), Some(o)) => (x, scale, o),
        _ => return false,
    };
    let bias = match tensor_dtype(node, 2) {
        Some(b) => b,
        None => return false,
    };
    if !is_mlx_float(x.dtype) || scale.dtype != x.dtype || bias != x.dtype || out.dtype != x.dtype {
        return false;
    }
    if x.shape.len() < 2 {
        return false;
    }
    let c = x.shape[1];
    if c <= 0 {
        return false;
    }
    for &d in &x.shape {
        if d <= 0 {
            return false; // need static dims to build the group reshape
        }
    }
    let groups = node.int_attr("num_groups", 0);
    if groups <= 0 || c % groups != 0 {
        return false;
    }
    // opset-21 per-channel scale/bias: shape [C].
    scale.shape.len() == 1 && scale.shape[0] == c
}

/// LpNormalization: single float input/output, p in {1,2}.
fn lp_norm_claim(node: &NodeView) -> bool {
    if node.num_inputs() != 1 || node.num_outputs() != 1 {
        return false;
    }
    let (x, out) = match (node.input_info(0), node.output_info(0)) {
        (Some(x), Some(o)) => (x, o),
        _ => return false,
    };
    if !is_mlx_float(x.dtype) || out.dtype != x.dtype || x.shape.is_empty() {
        return false;
    }
    let p = node.int_attr("p", 2);
    p == 1 || p == 2
}

/// BatchNormalization: inference (single-output) form, 5 float inputs sharing X's dtype.
fn batch_norm_claim(node: &NodeView) -> bool {
    if node.num_inputs() != 5 || node.num_outputs() != 1 {
        return false; // training outputs → CPU
    }
    let (x, out) = match (node.input_info(0), node.output_info(0)) {
        (Some(x), Some(o)) => (x, o),
        _ => return false,
    };
    if !is_mlx_float(x.dtype) || out.dtype != x.dtype || x.shape.len() < 2 {
        return false;
    }
    for i in 1..5 {
        match tensor_dtype(node, i) {
            Some(t) if t == x.dtype => {}
            _ => return false,
        }
    }
    node.int_attr("training_mode", 0) == 0
}

/// LRN (ai.onnx, across-channel): single float input/output of equal dtype, static shape with a
/// channel axis, and a valid window `size >= 1`.
fn lrn_claim(node: &NodeView) -> bool {
    if node.num_inputs() != 1 || node.num_outputs() != 1 {
        return false;
    }
    let (x, out) = match (node.input_info(0), node.output_info(0)) {
        (Some(x), Some(o)) => (x, o),
        _ => return false,
    };
    if !is_mlx_float(x.dtype) || out.dtype != x.dtype || x.shape.len() < 2 {
        return false;
    }
    for &d in &x.shape {
        if d <= 0 {
            return false; // need static dims to build the channel pad/slice
        }
    }
    node.int_attr("size", 0) >= 1
}

// ---- registration ------------------------------------------------------------------------------

#[allow(clippy::too_many_arguments)]
fn reg(
    registry: &mut OpRegistry,
    domain: &'static str,
    op_type: &'static str,
    min_opset: i32,
    handler: crate::registry::OpHandler,
    claim: crate::registry::ClaimPredicate,
) {
    registry.register(OpRegistration {
        domain,
        op_type,
        min_opset,
        max_opset: K_ANY_OPSET,
        handler,
        claim,
    });
}

pub fn register_norm(registry: &mut OpRegistry) {
    // RMSNormalization entered ai.onnx at opset 23.
    reg(registry, "", "RMSNormalization", 23, rms_norm_op, rms_norm_claim);
    // LayerNormalization entered ai.onnx at opset 17.
    reg(registry, "", "LayerNormalization", 17, layer_norm_op, layer_norm_claim);
    reg(registry, "", "GroupNormalization", K_ANY_OPSET, group_norm_op, group_norm_claim);
    reg(registry, "", "LpNormalization", K_ANY_OPSET, lp_norm_op, lp_norm_claim);
    reg(registry, "", "BatchNormalization", K_ANY_OPSET, batch_norm_op, batch_norm_claim);
    reg(registry, "", "LRN", K_ANY_OPSET, lrn_op, lrn_claim);
    reg(registry, "com.microsoft", "SimplifiedLayerNormalization", K_ANY_OPSET, simplified_layer_norm_op, simplified_layer_norm_claim);
    reg(registry, "com.microsoft", "SkipLayerNormalization", K_ANY_OPSET, skip_layer_norm_op, skip_layer_norm_claim);
    reg(registry, "com.microsoft", "SkipSimplifiedLayerNormalization", K_ANY_OPSET, skip_rms_norm_op, skip_rms_norm_claim);
}