onnxruntime-ep-mlx 0.29.4

//! State-space / KV-cache "misc" op handlers. Faithful port of the C++ `ops/ssm_misc.cc`:
//!   * TensorScatter (ai.onnx) — static-KV-cache scatter ("linear" mode) via slice_update /
//!     slice_update_dynamic.
//!   * CausalConvWithState (com.microsoft) — fused causal depthwise conv1d with carry state.
//!   * LinearAttention (com.microsoft) — chunked/recurrent linear attention (4 update rules,
//!     GQA) via static-length unrolling over the time axis T.
//!     Only statically translatable, MLX-supported forms are claimed; the rest fall to ORT CPU.

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

const T_INT64: ort::ONNXTensorElementDataType =
    ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64;
const I32: mlx::mlx_dtype = mlx::mlx_dtype__MLX_INT32;

// ---- shared helpers -----------------------------------------------------------------------------

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

fn str_attr(n: &NodeDesc, name: &str, dflt: &str) -> String {
    n.strings
        .get(name)
        .cloned()
        .unwrap_or_else(|| dflt.to_string())
}

fn norm_axis(axis: i64, rank: i32) -> i32 {
    let a = if axis < 0 { axis + rank as i64 } else { axis };
    a as i32
}

/// True when an optional input is omitted in the MIDDLE of the input list (an interior gap), which
/// the shared clustering pass cannot represent; such forms are left to ORT CPU.
fn has_interior_gap(node: &NodeView) -> bool {
    let n = node.num_inputs();
    let mut last_present = 0usize;
    let mut seen = false;
    for i in 0..n {
        if node.input_present(i) {
            last_present = i;
            seen = true;
        }
    }
    if !seen {
        return false;
    }
    (0..last_present).any(|i| !node.input_present(i))
}

// =============================================================================================
// TensorScatter (ai.onnx) — write `update` into `past_cache` along `axis` at write_indices[0]/0.
// =============================================================================================
fn tensor_scatter_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let past = ctx.resolve(&n.inputs[0])?;
    let update = ctx.resolve(&n.inputs[1])?;
    let rank = ctx.ndim(past) as i32;
    let axis = norm_axis(n.ints.get("axis").copied().unwrap_or(-2), rank);

    let present_arr = if present(n, 2) {
        // Decode form: dynamic offset from write_indices along `axis` (batch_size == 1 per claim).
        let wi = ctx.resolve(&n.inputs[2])?;
        let wi = ctx.astype(wi, I32)?;
        let axes = [axis];
        ctx.emit(|res, s| unsafe {
            mlx::mlx_slice_update_dynamic(res, past, update, wi, axes.as_ptr(), axes.len(), s)
        })?
    } else {
        // Prefill form: write the update block at offset 0 along `axis` for every batch.
        let start = vec![0i32; rank as usize];
        let mut stop = vec![0i32; rank as usize];
        for i in 0..rank {
            stop[i as usize] = ctx.dim(past, i);
        }
        stop[axis as usize] = ctx.dim(update, axis);
        let strides = vec![1i32; rank as usize];
        ctx.emit(|res, s| unsafe {
            mlx::mlx_slice_update(
                res,
                past,
                update,
                start.as_ptr(),
                start.len(),
                stop.as_ptr(),
                stop.len(),
                strides.as_ptr(),
                strides.len(),
                s,
            )
        })?
    };
    let cont = ctx.contiguous(present_arr)?;
    ctx.bind(&n.outputs[0], cont);
    Ok(())
}

fn tensor_scatter_claim(node: &NodeView) -> ClaimResult {
    let ni = node.num_inputs();
    require!(
        (ni == 2 || ni == 3) && node.num_outputs() == 1,
        "expects 2-3 inputs and 1 output, got {}in/{}out",
        ni,
        node.num_outputs()
    );
    let mode = node.string_attr("mode", "linear");
    require!(mode == "linear", "mode must be \"linear\", got {mode:?}");
    let (past, update, out) = match (node.input_info(0), node.input_info(1), node.output_info(0)) {
        (Some(a), Some(b), Some(c)) => (a, b, c),
        _ => deny!("missing tensor type/shape info on an input or the output"),
    };
    require!(
        is_mlx_float(past.dtype) && update.dtype == past.dtype && out.dtype == past.dtype,
        "past/update/output must share one float dtype, got {} / {} -> {}",
        crate::registry::ort_dtype_name(past.dtype),
        crate::registry::ort_dtype_name(update.dtype),
        crate::registry::ort_dtype_name(out.dtype)
    );
    if ni == 3 && node.input_present(2) {
        match node.input_info(2) {
            Some(wi) if wi.dtype == T_INT64 => {}
            Some(wi) => deny!(
                "write_indices must be int64, got {}",
                crate::registry::ort_dtype_name(wi.dtype)
            ),
            None => deny!("missing tensor type/shape info on write_indices"),
        }
        // Only batch_size == 1 is expressible as one dynamic slice.
        require!(
            !past.shape.is_empty() && past.shape[0] == 1,
            "dynamic TensorScatter requires past_cache batch size 1, got shape {:?}",
            past.shape
        );
    }
    Ok(())
}

// =============================================================================================
// CausalConvWithState (com.microsoft) — stateful causal depthwise conv1d.
// =============================================================================================
fn causal_conv_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let x = ctx.resolve(&n.inputs[0])?; // (B, C, L)
    let weight = ctx.resolve(&n.inputs[1])?; // (C, 1, k)
    let dt = ctx.dtype_of(x);
    let b = ctx.dim(x, 0);
    let c = ctx.dim(x, 1);
    let l = ctx.dim(x, 2);
    let k = ctx.dim(weight, 2);
    let state_window = *n.ints.get("state_window").unwrap_or(&0) as i32;

    let has_bias = present(n, 2);
    let has_state = present(n, 3);

    // Left context: past_state (B, C, k-1) or k-1 zeros. For k == 1 there is no carry state.
    let mut x_pad = x;
    if k > 1 {
        let state = if has_state {
            let state = ctx.resolve(&n.inputs[3])?;
            if state_window > 0 {
                let state = ctx.slice(
                    state,
                    &[state_window - 1, 0, 0, 0],
                    &[state_window, b, c, k - 1],
                )?;
                ctx.reshape(state, &[b, c, k - 1])?
            } else {
                state
            }
        } else {
            ctx.zeros(&[b, c, k - 1], dt)?
        };
        x_pad = ctx.concat2(state, x, 2)?; // (B, C, k-1+L)
    }

    // present_state = last k-1 columns of x_pad (boundary output -> contiguous).
    if n.outputs.len() >= 2 && !n.outputs[1].name.is_empty() {
        if state_window > 0 {
            if k > 1 && l > 0 {
                let kept = l.min(state_window);
                let leading = state_window - kept;
                let mut window = if leading > 0 {
                    Some(ctx.zeros(&[leading, b, c, k - 1], dt)?)
                } else {
                    None
                };
                for token in (l - kept)..l {
                    let carry = ctx.slice(x_pad, &[0, 0, token + 1], &[b, c, token + k])?;
                    let carry = ctx.expand_dims(carry, 0)?;
                    window = Some(match window {
                        Some(previous) => ctx.concat2(previous, carry, 0)?,
                        None => carry,
                    });
                }
                let ps = window
                    .ok_or_else(|| "MLX CausalConv state_window produced no states".to_string())?;
                let ps = ctx.contiguous(ps)?;
                ctx.bind(&n.outputs[1], ps);
            } else if has_state && l == 0 {
                let ps = ctx.resolve(&n.inputs[3])?;
                let ps = ctx.contiguous(ps)?;
                ctx.bind(&n.outputs[1], ps);
            } else {
                let z = ctx.zeros(&[state_window, b, c, k - 1], dt)?;
                ctx.bind(&n.outputs[1], z);
            }
        } else if k > 1 {
            let padded = ctx.dim(x_pad, 2);
            let ps = ctx.slice(x_pad, &[0, 0, padded - (k - 1)], &[b, c, padded])?;
            let ps = ctx.contiguous(ps)?;
            ctx.bind(&n.outputs[1], ps);
        } else {
            let z = ctx.zeros(&[b, c, 0], dt)?;
            ctx.bind(&n.outputs[1], z);
        }
    }

    // Depthwise conv1d: MLX uses NLC data and (C_out, kernel, C_in/groups) weights.
    let x_t = ctx.transpose(x_pad, &[0, 2, 1])?;
    let x_nlc = ctx.contiguous(x_t)?; // (B, k-1+L, C)
    let w_t = ctx.transpose(weight, &[0, 2, 1])?;
    let w_ckc = ctx.contiguous(w_t)?; // (C, k, 1)
    let y_nlc = ctx.emit(|res, s| unsafe { mlx::mlx_conv1d(res, x_nlc, w_ckc, 1, 0, 1, c, s) })?;
    let y_t = ctx.transpose(y_nlc, &[0, 2, 1])?;
    let mut y = ctx.contiguous(y_t)?; // (B, C, L)

    if has_bias {
        let bias = ctx.resolve(&n.inputs[2])?; // (C,)
        let b3 = ctx.reshape(bias, &[1, c, 1])?;
        y = ctx.add(y, b3)?;
    }

    let activation = str_attr(n, "activation", "none");
    if activation == "silu" || activation == "swish" {
        let sig = ctx.emit(|res, s| unsafe { mlx::mlx_sigmoid(res, y, s) })?;
        y = ctx.mul(y, sig)?;
    }
    ctx.bind(&n.outputs[0], y);
    Ok(())
}

fn causal_conv_claim(node: &NodeView) -> ClaimResult {
    let ni = node.num_inputs();
    require!((2..=4).contains(&ni), "expects 2-4 inputs, got {ni}");
    let no = node.num_outputs();
    require!((1..=2).contains(&no), "expects 1-2 outputs, got {no}");
    require!(
        !has_interior_gap(node),
        "optional inputs may only be omitted from the trailing end"
    );
    let (input, weight) = match (node.input_info(0), node.input_info(1)) {
        (Some(a), Some(b)) => (a, b),
        _ => deny!("missing tensor type/shape info on input or weight"),
    };
    require!(
        is_mlx_float(input.dtype) && weight.dtype == input.dtype,
        "input/weight must share one float dtype, got {} / {}",
        crate::registry::ort_dtype_name(input.dtype),
        crate::registry::ort_dtype_name(weight.dtype)
    );
    require!(
        input.shape.len() == 3 && weight.shape.len() == 3,
        "input and weight must both be rank 3, got ranks {} and {}",
        input.shape.len(),
        weight.shape.len()
    );
    let state_window = node.int_attr("state_window", 0);
    require!(
        (0..=8).contains(&state_window),
        "state_window must be in [0,8], got {state_window}"
    );
    let expected_state_shape = if state_window > 0 {
        vec![
            state_window,
            input.shape[0],
            input.shape[1],
            weight.shape[2] - 1,
        ]
    } else {
        vec![input.shape[0], input.shape[1], weight.shape[2] - 1]
    };
    if node.input_present(2) {
        match node.input_info(2) {
            Some(b) if b.dtype == input.dtype => {}
            Some(b) => deny!(
                "bias dtype must match input dtype {}, got {}",
                crate::registry::ort_dtype_name(input.dtype),
                crate::registry::ort_dtype_name(b.dtype)
            ),
            None => deny!("missing tensor type/shape info on bias"),
        }
    }
    if node.input_present(3) {
        match node.input_info(3) {
            Some(p) if p.dtype == input.dtype && p.shape == expected_state_shape => {}
            Some(p) => deny!(
                "past_state must have dtype {} and shape {:?}, got {} {:?}",
                crate::registry::ort_dtype_name(input.dtype),
                expected_state_shape,
                crate::registry::ort_dtype_name(p.dtype),
                p.shape
            ),
            None => deny!("missing tensor type/shape info on past_state"),
        }
    }
    if let Some(present) = node.output_info(1) {
        require!(
            present.dtype == input.dtype && present.shape == expected_state_shape,
            "present_state must have input dtype and shape {:?}",
            expected_state_shape
        );
    }
    let activation = node.string_attr("activation", "none");
    require!(
        activation == "none" || activation == "silu" || activation == "swish",
        "activation must be \"none\", \"silu\", or \"swish\", got {activation:?}"
    );
    Ok(())
}

// =============================================================================================
// LinearAttention (com.microsoft) — delta-rule linear attention, static-length unroll over T.
// =============================================================================================
fn rule_uses_decay(rule: &str) -> bool {
    rule == "gated" || rule == "gated_delta"
}
fn rule_uses_beta(rule: &str) -> bool {
    rule == "delta" || rule == "gated_delta"
}
fn is_known_rule(rule: &str) -> bool {
    matches!(rule, "linear" | "gated" | "delta" | "gated_delta")
}

fn la_scalar(
    ctx: &mut TranslationContext,
    value: f32,
    dt: mlx::mlx_dtype,
) -> Result<mlx::mlx_array, MlxError> {
    let s = ctx.scalar_f32(value);
    if dt == mlx::mlx_dtype__MLX_FLOAT32 {
        Ok(s)
    } else {
        ctx.astype(s, dt)
    }
}

/// From a (B, H, T, X) tensor pick time-step `t` as a (B, H, X) slab.
fn time_slab(
    ctx: &mut TranslationContext,
    a: mlx::mlx_array,
    t: i32,
    b: i32,
    h: i32,
    x: i32,
) -> Result<mlx::mlx_array, MlxError> {
    let s = ctx.slice(a, &[0, 0, t, 0], &[b, h, t + 1, x])?;
    ctx.reshape(s, &[b, h, x])
}

/// From a (B, H, T) tensor pick time-step `t` as a (B, H) slab.
fn time_slab2(
    ctx: &mut TranslationContext,
    a: mlx::mlx_array,
    t: i32,
    b: i32,
    h: i32,
) -> Result<mlx::mlx_array, MlxError> {
    let s = ctx.slice(a, &[0, 0, t], &[b, h, t + 1])?;
    ctx.reshape(s, &[b, h])
}

fn repeat_axis(
    ctx: &mut TranslationContext,
    a: mlx::mlx_array,
    repeats: i32,
    axis: i32,
) -> Result<mlx::mlx_array, MlxError> {
    if repeats == 1 {
        return Ok(a);
    }
    ctx.emit(|res, s| unsafe { mlx::mlx_repeat_axis(res, a, repeats, axis, s) })
}

// ---- chunked gated-delta helpers ----------------------------------------------------------------

const CHUNK_SIZE: i32 = 64;

/// Zero-pad `a` at the END of `axis` by `high` elements, then force contiguous (chunk reshapes and
/// the batched matmuls below all require a dense buffer). `high == 0` still returns a contiguous
/// copy so a subsequent `reshape` never sees a strided view.
fn pad_axis_end(
    ctx: &mut TranslationContext,
    a: mlx::mlx_array,
    axis: i32,
    high: i32,
    dt: mlx::mlx_dtype,
) -> Result<mlx::mlx_array, MlxError> {
    if high == 0 {
        return ctx.contiguous(a);
    }
    let zero = la_scalar(ctx, 0.0, dt)?;
    let axes = [axis];
    let lo = [0i32];
    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,
            zero,
            mode.as_ptr() as *const c_char,
            s,
        )
    })?;
    ctx.contiguous(out)
}

fn tril(
    ctx: &mut TranslationContext,
    a: mlx::mlx_array,
    k: i32,
) -> Result<mlx::mlx_array, MlxError> {
    ctx.emit(|res, s| unsafe { mlx::mlx_tril(res, a, k, s) })
}

fn exp_arr(ctx: &mut TranslationContext, a: mlx::mlx_array) -> Result<mlx::mlx_array, MlxError> {
    ctx.emit(|res, s| unsafe { mlx::mlx_exp(res, a, s) })
}

/// Transpose the last two axes of a rank-N array (N in {4,5}).
fn transpose_last2(
    ctx: &mut TranslationContext,
    a: mlx::mlx_array,
    rank: i32,
) -> Result<mlx::mlx_array, MlxError> {
    let mut axes: Vec<i32> = (0..rank).collect();
    axes.swap((rank - 1) as usize, (rank - 2) as usize);
    ctx.transpose(a, &axes)
}

/// Select chunk `i` along axis 2 of a rank-5 `[B,H,nc,C,X]` tensor, returning `[B,H,C,X]`.
fn chunk5(
    ctx: &mut TranslationContext,
    a: mlx::mlx_array,
    i: i32,
    b: i32,
    h: i32,
    c: i32,
    x: i32,
) -> Result<mlx::mlx_array, MlxError> {
    let s = ctx.slice(a, &[0, 0, i, 0, 0], &[b, h, i + 1, c, x])?;
    ctx.reshape(s, &[b, h, c, x])
}

/// Select chunk `i` along axis 2 of a rank-4 `[B,H,nc,C]` tensor, returning `[B,H,C]`.
fn chunk4(
    ctx: &mut TranslationContext,
    a: mlx::mlx_array,
    i: i32,
    b: i32,
    h: i32,
    c: i32,
) -> Result<mlx::mlx_array, MlxError> {
    let s = ctx.slice(a, &[0, 0, i, 0], &[b, h, i + 1, c])?;
    ctx.reshape(s, &[b, h, c])
}

/// Chunked gated-delta / delta linear attention (chunk_size = 64). A faithful MLX port of HF
/// Qwen3.5's `torch_chunk_gated_delta_rule`; numerically identical to the recurrent unroll (the
/// ground-truth handler), but expressed as batched matmuls so PREFILL fills the GPU. All ops are
/// static-shape (the C-axis forward-substitution is a fixed 64-iteration loop), so the whole thing
/// is captured inside `mlx_compile`.
///
/// Inputs are already in `[B,H,T,·]` head layout with GQA tiling applied and the query pre-scaled.
/// `g_bht` is the per-head SCALAR decay `[B,H,T]` (zeros for the `delta` rule); `beta_bht` is
/// `[B,H,T]`. `state` is the initial `[B,H,d_k,d_v]` recurrent state.
#[allow(clippy::too_many_arguments)]
fn linear_attention_chunked(
    ctx: &mut TranslationContext,
    n: &NodeDesc,
    q4: mlx::mlx_array,
    k4: mlx::mlx_array,
    v4: mlx::mlx_array,
    g_bht: mlx::mlx_array,
    beta_bht: mlx::mlx_array,
    mut state: mlx::mlx_array,
    b: i32,
    h: i32,
    d_k: i32,
    d_v: i32,
    t_len: i32,
    dt: mlx::mlx_dtype,
) -> Result<(), MlxError> {
    let c = CHUNK_SIZE;
    let pad = (c - t_len % c) % c;
    let tp = t_len + pad;
    let nc = tp / c;

    // Pad time to a chunk multiple (query/key/value with zeros; beta/g with zeros too).
    let q = pad_axis_end(ctx, q4, 2, pad, dt)?; // [B,H,tp,d_k]
    let k = pad_axis_end(ctx, k4, 2, pad, dt)?; // [B,H,tp,d_k]
    let v = pad_axis_end(ctx, v4, 2, pad, dt)?; // [B,H,tp,d_v]
    let beta = pad_axis_end(ctx, beta_bht, 2, pad, dt)?; // [B,H,tp]
    let g = pad_axis_end(ctx, g_bht, 2, pad, dt)?; // [B,H,tp]

    // v_beta / k_beta on the flat (unchunked) tensors, then reshape everything into chunks.
    let beta_col = ctx.expand_dims(beta, 3)?; // [B,H,tp,1]
    let v_beta = ctx.mul(v, beta_col)?; // [B,H,tp,d_v]
    let k_beta = ctx.mul(k, beta_col)?; // [B,H,tp,d_k]

    let q_c = ctx.reshape(q, &[b, h, nc, c, d_k])?;
    let k_c = ctx.reshape(k, &[b, h, nc, c, d_k])?;
    let k_beta_c = ctx.reshape(k_beta, &[b, h, nc, c, d_k])?;
    let v_beta_c = ctx.reshape(v_beta, &[b, h, nc, c, d_v])?;
    let g_c = ctx.reshape(g, &[b, h, nc, c])?; // [B,H,nc,C]

    // Cumulative decay within each chunk and the pairwise decay matrix.
    let g_cum = ctx.emit(|res, s| unsafe { mlx::mlx_cumsum(res, g_c, 3, false, true, s) })?; // [B,H,nc,C]
    let gi = ctx.expand_dims(g_cum, 4)?; // [B,H,nc,C,1]
    let gj = ctx.expand_dims(g_cum, 3)?; // [B,H,nc,1,C]
    let diff = ctx.sub(gi, gj)?; // [B,H,nc,C,C]
    let low = tril(ctx, diff, 0)?;
    let dexp = exp_arr(ctx, low)?;
    let decay_mask = tril(ctx, dexp, 0)?; // [B,H,nc,C,C]

    // Upper-triangular (incl. diagonal) boolean mask, broadcast over batch/chunk.
    let ones_cc = ctx.emit(|res, s| unsafe {
        mlx::mlx_ones(res, [c, c].as_ptr(), 2, mlx::mlx_dtype__MLX_BOOL, s)
    })?;
    let mask = ctx.emit(|res, s| unsafe { mlx::mlx_triu(res, ones_cc, 0, s) })?;
    let zero = la_scalar(ctx, 0.0, dt)?;

    // attn = -((k_beta @ key^T) * decay_mask) with the upper triangle (incl. diagonal) zeroed.
    let k_t = transpose_last2(ctx, k_c, 5)?; // [B,H,nc,d_k,C]
    let kk = ctx.matmul(k_beta_c, k_t)?; // [B,H,nc,C,C]
    let kk = ctx.mul(kk, decay_mask)?;
    let neg = ctx.emit(|res, s| unsafe { mlx::mlx_negative(res, kk, s) })?;
    let mut attn = ctx.where_(mask, zero, neg)?; // [B,H,nc,C,C]

    // Forward substitution: build (I - A)^-1 in place (fixed C-iteration loop, static shapes).
    for i in 1..c {
        let row = ctx.slice(attn, &[0, 0, 0, i, 0], &[b, h, nc, i + 1, i])?; // [B,H,nc,1,i]
        let sub = ctx.slice(attn, &[0, 0, 0, 0, 0], &[b, h, nc, i, i])?; // [B,H,nc,i,i]
        let prod = ctx.matmul(row, sub)?; // [B,H,nc,1,i]
        let new_row = ctx.add(row, prod)?;
        attn = ctx.slice_update(attn, new_row, &[0, 0, 0, i, 0], &[b, h, nc, i + 1, i])?;
    }
    let eye = ctx.emit(|res, s| unsafe { mlx::mlx_eye(res, c, c, 0, dt, s) })?;
    let attn = ctx.add(attn, eye)?; // [B,H,nc,C,C]

    let u_val = ctx.matmul(attn, v_beta_c)?; // [B,H,nc,C,d_v] ("pseudo-values")
    let g_exp = exp_arr(ctx, g_cum)?; // [B,H,nc,C]
    let g_exp_col = ctx.expand_dims(g_exp, 4)?; // [B,H,nc,C,1]
    let k_scaled = ctx.mul(k_beta_c, g_exp_col)?; // [B,H,nc,C,d_k]
    let k_cumdecay = ctx.matmul(attn, k_scaled)?; // [B,H,nc,C,d_k]

    // Sequential scan over chunks; each step is a handful of batched matmuls over [B,H,C,·].
    let mut out_chunks: Vec<mlx::mlx_array> = Vec::with_capacity(nc as usize);
    for i in 0..nc {
        let q_i = chunk5(ctx, q_c, i, b, h, c, d_k)?; // [B,H,C,d_k]
        let k_i = chunk5(ctx, k_c, i, b, h, c, d_k)?; // [B,H,C,d_k]
        let v_i = chunk5(ctx, u_val, i, b, h, c, d_v)?; // [B,H,C,d_v]
        let dm_i = chunk5(ctx, decay_mask, i, b, h, c, c)?; // [B,H,C,C]
        let gc_i = chunk4(ctx, g_cum, i, b, h, c)?; // [B,H,C]
        let kcd_i = chunk5(ctx, k_cumdecay, i, b, h, c, d_k)?; // [B,H,C,d_k]

        let k_i_t = transpose_last2(ctx, k_i, 4)?; // [B,H,d_k,C]
        let attn_i = ctx.matmul(q_i, k_i_t)?; // [B,H,C,C]
        let attn_i = ctx.mul(attn_i, dm_i)?;

        let v_prime = ctx.matmul(kcd_i, state)?; // [B,H,C,d_v]
        let v_new = ctx.sub(v_i, v_prime)?; // [B,H,C,d_v]

        let gc_i_exp = exp_arr(ctx, gc_i)?; // [B,H,C]
        let gc_i_col = ctx.expand_dims(gc_i_exp, 3)?; // [B,H,C,1]
        let q_dec = ctx.mul(q_i, gc_i_col)?; // [B,H,C,d_k]
        let attn_inter = ctx.matmul(q_dec, state)?; // [B,H,C,d_v]
        let intra = ctx.matmul(attn_i, v_new)?; // [B,H,C,d_v]
        let out_i = ctx.add(attn_inter, intra)?; // [B,H,C,d_v]
        out_chunks.push(out_i);

        // State update: decay carried state and add the chunk's key/value outer product.
        let g_last = ctx.slice(gc_i, &[0, 0, c - 1], &[b, h, c])?; // [B,H,1]
        let g_last_exp = exp_arr(ctx, g_last)?; // [B,H,1]
        let g_last_e = ctx.expand_dims(g_last_exp, 3)?; // [B,H,1,1]
        let decayed = ctx.mul(state, g_last_e)?; // [B,H,d_k,d_v]

        let w = ctx.sub(g_last, gc_i)?; // [B,H,C]  (g_last - g_cum)
        let w_exp = exp_arr(ctx, w)?;
        let w_col = ctx.expand_dims(w_exp, 3)?; // [B,H,C,1]
        let kw = ctx.mul(k_i, w_col)?; // [B,H,C,d_k]
        let kw_t = transpose_last2(ctx, kw, 4)?; // [B,H,d_k,C]
        let upd = ctx.matmul(kw_t, v_new)?; // [B,H,d_k,d_v]
        state = ctx.add(decayed, upd)?;
    }

    if !n.outputs.is_empty() && !n.outputs[0].name.is_empty() {
        let out5 = ctx.stack(&out_chunks, 2)?; // [B,H,nc,C,d_v]
        let out_flat = ctx.reshape(out5, &[b, h, tp, d_v])?;
        let out = ctx.slice(out_flat, &[0, 0, 0, 0], &[b, h, t_len, d_v])?; // drop padding
        let out = ctx.transpose(out, &[0, 2, 1, 3])?; // [B,T,H,d_v]
        let out = ctx.reshape(out, &[b, t_len, h * d_v])?;
        let out = ctx.contiguous(out)?;
        ctx.bind(&n.outputs[0], out);
    }
    if n.outputs.len() >= 2 && !n.outputs[1].name.is_empty() {
        let ps = ctx.contiguous(state)?;
        ctx.bind(&n.outputs[1], ps);
    }
    Ok(())
}

fn linear_attention_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let rule = str_attr(n, "update_rule", "gated_delta");
    let uses_decay = rule_uses_decay(&rule);
    let uses_beta = rule_uses_beta(&rule);
    let hq = *n
        .ints
        .get("q_num_heads")
        .ok_or("MLX LinearAttention: q_num_heads missing")? as i32;
    let h = *n
        .ints
        .get("kv_num_heads")
        .ok_or("MLX LinearAttention: kv_num_heads missing")? as i32;
    let state_window = *n.ints.get("state_window").unwrap_or(&0) as i32;

    let query = ctx.resolve(&n.inputs[0])?; // (B, T, Hq*d_k)
    let key = ctx.resolve(&n.inputs[1])?; // (B, T, Hk*d_k)
    let value = ctx.resolve(&n.inputs[2])?; // (B, T, H*d_v)
    let dt = ctx.dtype_of(query);
    let state_dt = if present(n, 3) {
        let past = ctx.resolve(&n.inputs[3])?;
        ctx.dtype_of(past)
    } else if n.outputs.len() >= 2 && !n.outputs[1].name.is_empty() {
        mlx_dtype_from_onnx(n.outputs[1].otype)
    } else {
        dt
    };

    let qsh = ctx.shape_of(query);
    let vsh = ctx.shape_of(value);
    let b = qsh[0];
    let t_len = qsh[1];
    // A dynamic (symbolic) time extent only survives into a SHAPELESS trace; the static-length
    // unroll below needs a concrete T. Fail the trace so the plan falls back to the eager translator
    // (or the shape-keyed general trace), where `shape_of` reports the real extent. The claim admits
    // dynamic-time nodes precisely because this guard keeps them correct.
    if t_len < 0 {
        return Err(
            "MLX LinearAttention: dynamic time extent in a shapeless trace — falls back to eager"
                .to_string(),
        );
    }
    let d_k = qsh[2] / hq;
    let d_v = vsh[2] / h;
    let n_k_heads = ctx.shape_of(key)[2] / d_k;
    if n_k_heads <= 0 || h % n_k_heads != 0 {
        return Err("MLX LinearAttention: kv_num_heads must be divisible by key heads".to_string());
    }
    let output_heads = hq.max(h);

    let scale_attr = n.floats.get("scale").copied().unwrap_or(0.0);
    let scale = if scale_attr != 0.0 {
        scale_attr
    } else {
        1.0 / (d_k as f32).sqrt()
    };

    let has_past = present(n, 3);
    let mut state = if has_past {
        let past = ctx.resolve(&n.inputs[3])?;
        if state_window > 0 {
            let past = ctx.slice(
                past,
                &[state_window - 1, 0, 0, 0, 0],
                &[state_window, b, h, d_k, d_v],
            )?;
            ctx.reshape(past, &[b, h, d_k, d_v])?
        } else {
            past
        }
    } else {
        ctx.zeros(&[b, h, d_k, d_v], state_dt)?
    };

    // Zero-length time axis: no steps run. output empty; present_state == state.
    if t_len == 0 {
        if !n.outputs.is_empty() && !n.outputs[0].name.is_empty() {
            let z = ctx.zeros(&[b, 0, output_heads * d_v], dt)?;
            ctx.bind(&n.outputs[0], z);
        }
        if n.outputs.len() >= 2 && !n.outputs[1].name.is_empty() {
            let ps = if state_window > 0 {
                if has_past {
                    ctx.resolve(&n.inputs[3])?
                } else {
                    ctx.zeros(&[state_window, b, h, d_k, d_v], state_dt)?
                }
            } else {
                state
            };
            let ps = ctx.contiguous(ps)?;
            ctx.bind(&n.outputs[1], ps);
        }
        return Ok(());
    }

    // 3D -> 4D (B, H, T, ·): reshape by head count, transpose heads before T, tile Q/K for GQA.
    let to_heads = |ctx: &mut TranslationContext,
                    a: mlx::mlx_array,
                    heads: i32,
                    last: i32|
     -> Result<mlx::mlx_array, MlxError> {
        let r = ctx.reshape(a, &[b, t_len, heads, last])?;
        ctx.transpose(r, &[0, 2, 1, 3])
    };
    let query = ctx.astype(query, state_dt)?;
    let key = ctx.astype(key, state_dt)?;
    let value = ctx.astype(value, state_dt)?;
    let q_heads = to_heads(ctx, query, hq, d_k)?;
    let q4 = if hq < h {
        repeat_axis(ctx, q_heads, h / hq, 1)? // inverse GQA: one Q head serves multiple KV states
    } else {
        q_heads
    };
    let k_heads = to_heads(ctx, key, n_k_heads, d_k)?;
    let k4 = repeat_axis(ctx, k_heads, h / n_k_heads, 1)?; // each K head serves one or more KV states
    let v4 = to_heads(ctx, value, h, d_v)?; // (B, H, T, d_v)
    let scale_s = la_scalar(ctx, scale, state_dt)?;
    let q4 = ctx.mul(q4, scale_s)?; // scaled query

    // Decay may be per-head-per-key-dim `[B,T,H*d_k]` (the op-test / ORT reference form) OR a
    // per-head SCALAR `[B,T,H]` (Qwen3.5 GatedDeltaNet exports one gate per value head). Detect it
    // from the input's last dim; the scalar form broadcasts across (d_k, d_v).
    let decay_per_head_scalar = uses_decay && {
        let d = ctx.resolve(&n.inputs[4])?;
        let ds = ctx.shape_of(d);
        ds.len() == 3 && ds[2] == h
    };
    let decay4 = if !uses_decay {
        None
    } else if decay_per_head_scalar {
        // [B,T,H] -> [B,H,T]
        let d = ctx.resolve(&n.inputs[4])?;
        let d = ctx.astype(d, state_dt)?;
        Some(ctx.transpose(d, &[0, 2, 1])?)
    } else {
        let d = ctx.resolve(&n.inputs[4])?;
        let d = ctx.astype(d, state_dt)?;
        Some(to_heads(ctx, d, h, d_k)?)
    };
    let beta3 = if uses_beta {
        let bta = ctx.resolve(&n.inputs[5])?;
        let bta = ctx.astype(bta, state_dt)?;
        let beta = ctx.transpose(bta, &[0, 2, 1])?;
        Some(if ctx.shape_of(beta)[1] == 1 && h > 1 {
            repeat_axis(ctx, beta, h, 1)?
        } else {
            beta
        })
    } else {
        None
    };

    // Prefill fast path: the CHUNKED gated-delta algorithm (chunk_size = 64) is numerically
    // identical (in exact arithmetic) to the recurrent unroll but expressed as batched matmuls, so
    // it fills the GPU on long sequences. It is used ONLY for `gated_delta` with a per-head SCALAR
    // decay layout (the Qwen3.5 GatedDeltaNet form). The `delta` rule is deliberately excluded: with
    // no decay to bound it, the `(I - A)^-1` forward substitution amplifies fp32 rounding and
    // diverges from ORT CPU (verified: ~1e-1..1e1 abs error at T>=128), so it stays on the stable
    // recurrent path — as do `linear` / `gated` (no delta correction) and the per-key-dim decay
    // layout. Decode / small T (t_len <= chunk_size) also stay recurrent (best for T == 1).
    let use_chunked = rule == "gated_delta"
        && decay_per_head_scalar
        && t_len > CHUNK_SIZE
        && hq == h
        && n_k_heads == h
        && state_dt == dt
        && state_window == 0;
    if use_chunked {
        let g_bht = decay4.expect("gated_delta carries decay"); // [B,H,T] scalar decay
        let beta_bht = beta3.expect("gated_delta carries beta"); // [B,H,T]
        return linear_attention_chunked(
            ctx, n, q4, k4, v4, g_bht, beta_bht, state, b, h, d_k, d_v, t_len, dt,
        );
    }

    let mut outs: Vec<mlx::mlx_array> = Vec::with_capacity(t_len as usize);
    let mut state_history: Vec<mlx::mlx_array> = if state_window > 0 {
        Vec::with_capacity(state_window as usize)
    } else {
        Vec::new()
    };
    for t in 0..t_len {
        if let Some(decay4) = decay4 {
            let g = if decay_per_head_scalar {
                let slab = time_slab2(ctx, decay4, t, b, h)?; // (B, H)
                let g = ctx.emit(|res, s| unsafe { mlx::mlx_exp(res, slab, s) })?;
                let g = ctx.expand_dims(g, 2)?; // (B,H,1)
                ctx.expand_dims(g, 3)? // (B,H,1,1) — broadcasts over (d_k,d_v)
            } else {
                let slab = time_slab(ctx, decay4, t, b, h, d_k)?; // (B, H, d_k)
                let g = ctx.emit(|res, s| unsafe { mlx::mlx_exp(res, slab, s) })?;
                ctx.expand_dims(g, 3)? // (B,H,d_k,1)
            };
            state = ctx.mul(state, g)?;
        }
        let k_t = time_slab(ctx, k4, t, b, h, d_k)?; // (B, H, d_k)
        // retrieval = squeeze(k_row @ state)
        let k_row = ctx.expand_dims(k_t, 2)?; // (B,H,1,d_k)
        let retrieval_m = ctx.matmul(k_row, state)?; // (B,H,1,d_v)
        let retrieval = ctx.squeeze(retrieval_m, 2)?; // (B,H,d_v)

        let v_t = time_slab(ctx, v4, t, b, h, d_v)?; // (B, H, d_v)
        let delta = if let Some(beta3) = beta3 {
            let beta_t = time_slab2(ctx, beta3, t, b, h)?; // (B, H)
            let diff = ctx.sub(v_t, retrieval)?;
            let beta_e = ctx.expand_dims(beta_t, 2)?; // (B,H,1)
            ctx.mul(diff, beta_e)?
        } else {
            v_t
        };
        // outer = k_col @ delta_row: (B,H,d_k,1) @ (B,H,1,d_v) -> (B,H,d_k,d_v)
        let k_col = ctx.expand_dims(k_t, 3)?;
        let delta_row = ctx.expand_dims(delta, 2)?;
        let outer = ctx.matmul(k_col, delta_row)?;
        state = ctx.add(state, outer)?;
        if state_window > 0 {
            if state_history.len() == state_window as usize {
                state_history.remove(0);
            }
            state_history.push(state);
        }

        let readout_state = if hq > h {
            repeat_axis(ctx, state, hq / h, 1)?
        } else {
            state
        };
        let q_t = time_slab(ctx, q4, t, b, output_heads, d_k)?;
        let q_row = ctx.expand_dims(q_t, 2)?; // (B,H,1,d_k)
        let out_m = ctx.matmul(q_row, readout_state)?; // (B,Hout,1,d_v)
        let out_t = ctx.squeeze(out_m, 2)?; // (B,H,d_v)
        outs.push(out_t);
    }

    if !n.outputs.is_empty() && !n.outputs[0].name.is_empty() {
        // Assemble output (B, T, H*d_v): each step's (B, H, d_v) reshapes to (B, 1, H*d_v).
        let mut out = ctx.reshape(outs[0], &[b, 1, output_heads * d_v])?;
        for &step_out in outs.iter().skip(1) {
            let slab = ctx.reshape(step_out, &[b, 1, output_heads * d_v])?;
            out = ctx.concat2(out, slab, 1)?;
        }
        let out = ctx.contiguous(out)?;
        let out = ctx.astype(out, dt)?;
        ctx.bind(&n.outputs[0], out);
    }
    if n.outputs.len() >= 2 && !n.outputs[1].name.is_empty() {
        let ps = if state_window > 0 {
            let kept = t_len.min(state_window);
            let leading = state_window - kept;
            let mut window = if leading > 0 {
                Some(ctx.zeros(&[leading, b, h, d_k, d_v], state_dt)?)
            } else {
                None
            };
            for &step_state in &state_history {
                let slot = ctx.expand_dims(step_state, 0)?;
                window = Some(match window {
                    Some(previous) => ctx.concat2(previous, slot, 0)?,
                    None => slot,
                });
            }
            window
                .ok_or_else(|| "MLX LinearAttention state_window produced no states".to_string())?
        } else {
            state
        };
        let ps = ctx.contiguous(ps)?;
        ctx.bind(&n.outputs[1], ps);
    }
    Ok(())
}

fn linear_attention_common_claim(node: &NodeView, allow_mixed_state: bool) -> ClaimResult {
    require!(
        node.num_inputs() >= 3 && node.num_outputs() >= 1,
        "expects at least 3 inputs and 1 output, got {}in/{}out",
        node.num_inputs(),
        node.num_outputs()
    );
    let state_window = node.int_attr("state_window", 0);
    require!(
        (0..=8).contains(&state_window),
        "state_window must be in [0,8], got {state_window}"
    );
    let rule = node.string_attr("update_rule", "gated_delta");
    require!(is_known_rule(&rule), "unsupported update_rule {rule:?}");
    let hq = node.int_attr("q_num_heads", 0);
    let h = node.int_attr("kv_num_heads", 0);
    require!(
        hq > 0 && h > 0 && (hq % h == 0 || h % hq == 0),
        "q_num_heads and kv_num_heads must be positive and one must divide the other, got q={hq}, kv={h}"
    );
    let (q, k, v) = match (node.input_info(0), node.input_info(1), node.input_info(2)) {
        (Some(a), Some(b), Some(c)) => (a, b, c),
        _ => deny!("missing tensor type/shape info on query, key, or value"),
    };
    require!(
        is_mlx_float(q.dtype) && k.dtype == q.dtype && v.dtype == q.dtype,
        "query/key/value must share one float dtype, got {} / {} / {}",
        crate::registry::ort_dtype_name(q.dtype),
        crate::registry::ort_dtype_name(k.dtype),
        crate::registry::ort_dtype_name(v.dtype)
    );
    require!(
        q.shape.len() == 3,
        "query must be rank 3, got shape {:?}",
        q.shape
    );
    require!(
        q.shape[2] > 0 && q.shape[2] % hq == 0,
        "query hidden size must be a positive multiple of q_num_heads"
    );
    let d_k = q.shape[2] / hq;
    require!(
        k.shape.len() == 3
            && k.shape[2] > 0
            && k.shape[2] % d_k == 0
            && h % (k.shape[2] / d_k) == 0,
        "key heads inferred from key hidden size must divide kv_num_heads"
    );
    require!(
        v.shape.len() == 3 && v.shape[2] > 0 && v.shape[2] % h == 0,
        "value hidden size must be a positive multiple of kv_num_heads"
    );
    let d_v = v.shape[2] / h;
    let expected_state_shape = if state_window > 0 {
        vec![state_window, q.shape[0], h, d_k, d_v]
    } else {
        vec![q.shape[0], h, d_k, d_v]
    };
    if let Some(past) = node.input_info(3) {
        require!(
            past.shape == expected_state_shape,
            "past_state shape {:?} must be {:?}",
            past.shape,
            expected_state_shape
        );
    }
    if let Some(present_state) = node.output_info(1) {
        require!(
            present_state.shape == expected_state_shape,
            "present_state shape {:?} must be {:?}",
            present_state.shape,
            expected_state_shape
        );
    }
    // The time dimension may be DYNAMIC (symbolic `[-1]`, as in real Qwen3.5/Qwen3-Next decoder
    // exports): the shape-keyed general trace resolves the concrete extent at trace time, and the
    // handler unrolls over it (with a guard that falls back to eager if it is ever traced shapeless).
    let float_ok = |i: usize| -> bool {
        if !node.input_present(i) {
            return true;
        }
        matches!(node.input_info(i), Some(info) if info.dtype == q.dtype)
    };
    if allow_mixed_state {
        if node.input_present(3) {
            require!(
                matches!(node.input_info(3), Some(info) if is_mlx_float(info.dtype)),
                "past_state must have an MLX float dtype"
            );
        }
    } else {
        require!(
            float_ok(3),
            "past_state dtype must match query dtype {}",
            crate::registry::ort_dtype_name(q.dtype)
        );
    }
    require!(
        float_ok(4),
        "decay dtype must match query dtype {}",
        crate::registry::ort_dtype_name(q.dtype)
    );
    require!(
        float_ok(5),
        "beta dtype must match query dtype {}",
        crate::registry::ort_dtype_name(q.dtype)
    );
    require!(
        !rule_uses_decay(&rule) || node.input_present(4),
        "update_rule {rule:?} requires the decay input"
    );
    require!(
        !rule_uses_beta(&rule) || node.input_present(5),
        "update_rule {rule:?} requires the beta input"
    );
    if node.input_present(5) {
        let beta = node.input_info(5).unwrap();
        require!(
            beta.shape.len() == 3
                && beta.shape[0] == q.shape[0]
                && beta.shape[1] == q.shape[1]
                && (beta.shape[2] == 1 || beta.shape[2] == h),
            "beta must have shape [B,T,1] or [B,T,kv_num_heads]"
        );
    }
    Ok(())
}

fn linear_attention_claim(node: &NodeView) -> ClaimResult {
    linear_attention_common_claim(node, false)
}

fn linear_attention_standard_claim(node: &NodeView) -> ClaimResult {
    linear_attention_common_claim(node, true)?;
    let q = node.input_info(0).unwrap();
    let k = node.input_info(1).unwrap();
    let v = node.input_info(2).unwrap();
    let out = node.output_info(0).unwrap();
    let hq = node.int_attr("q_num_heads", 0);
    let h = node.int_attr("kv_num_heads", 0);
    require!(
        hq % h == 0,
        "ai.onnx LinearAttention requires q_num_heads divisible by kv_num_heads"
    );
    let d_k = q.shape[2] / hq;
    let d_v = v.shape[2] / h;
    require!(
        k.shape[2] == h * d_k,
        "ai.onnx LinearAttention key must use kv_num_heads heads"
    );
    require!(
        out.shape.len() == 3 && out.shape[2] == hq * d_v,
        "ai.onnx LinearAttention output must have hidden size q_num_heads*d_v"
    );
    if let Some(state_out) = node.output_info(1) {
        require!(
            is_mlx_float(state_out.dtype),
            "ai.onnx LinearAttention present_state must have an MLX float dtype"
        );
        if let Some(state_in) = node.input_info(3) {
            require!(
                state_out.dtype == state_in.dtype,
                "present_state dtype must match past_state dtype"
            );
        }
    }
    Ok(())
}

// ---- ORT 1.29 fused gates ----------------------------------------------------------------------

fn gated_add_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let x = ctx.resolve(&n.inputs[0])?;
    let y = ctx.resolve(&n.inputs[1])?;
    let gate = ctx.resolve(&n.inputs[2])?;
    let scaled = ctx.mul(y, gate)?;
    let out = ctx.add(x, scaled)?;
    ctx.bind(&n.outputs[0], out);
    Ok(())
}

fn gated_add_claim(node: &NodeView) -> ClaimResult {
    require!(
        node.num_inputs() == 3 && node.num_outputs() == 1,
        "GatedAdd expects 3 inputs and 1 output"
    );
    let (x, y, gate, out) = match (
        node.input_info(0),
        node.input_info(1),
        node.input_info(2),
        node.output_info(0),
    ) {
        (Some(x), Some(y), Some(gate), Some(out)) => (x, y, gate, out),
        _ => deny!("GatedAdd requires typed input/output tensors"),
    };
    require!(
        is_mlx_float(x.dtype)
            && y.dtype == x.dtype
            && gate.dtype == x.dtype
            && out.dtype == x.dtype,
        "X, Y, gate, and output must share one MLX float dtype"
    );
    require!(
        !x.shape.is_empty() && x.shape.last().is_some_and(|&d| d > 0),
        "X must have rank >= 1 and a positive static last dimension"
    );
    require!(
        y.shape == x.shape && out.shape == x.shape,
        "Y and output must have the same shape as X"
    );
    let mut gate_shape = x.shape.clone();
    *gate_shape.last_mut().unwrap() = 1;
    require!(
        gate.shape == gate_shape,
        "gate shape must equal X shape with the last dimension replaced by 1"
    );
    Ok(())
}

fn linear_attention_gate_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let a = ctx.resolve(&n.inputs[0])?;
    let out_dt = ctx.dtype_of(a);
    let a32 = ctx.astype(a, mlx::mlx_dtype__MLX_FLOAT32)?;
    let bias = ctx.resolve(&n.inputs[1])?;
    let shifted = ctx.add(a32, bias)?;
    let zero = ctx.zeros_like(shifted)?;
    let softplus = ctx.binary(mlx::mlx_logaddexp, zero, shifted)?;
    let scale = ctx.resolve(&n.inputs[2])?;
    let decay32 = ctx.mul(softplus, scale)?;
    let decay = ctx.astype(decay32, out_dt)?;
    ctx.bind(&n.outputs[0], decay);

    if n.outputs.len() > 1 && !n.outputs[1].name.is_empty() {
        let b = ctx.resolve(&n.inputs[3])?;
        let b32 = ctx.astype(b, mlx::mlx_dtype__MLX_FLOAT32)?;
        let beta32 = ctx.unary(mlx::mlx_sigmoid, b32)?;
        let beta = ctx.astype(beta32, out_dt)?;
        ctx.bind(&n.outputs[1], beta);
    }
    Ok(())
}

fn linear_attention_gate_claim(node: &NodeView) -> ClaimResult {
    require!(
        node.num_inputs() == 4 && (1..=2).contains(&node.num_outputs()),
        "LinearAttentionGate expects 4 input slots and 1-2 outputs"
    );
    let a = match node.input_info(0) {
        Some(a) if is_mlx_float(a.dtype) => a,
        Some(_) => deny!("a must have an MLX float dtype"),
        None => deny!("a lacks tensor type/shape info"),
    };
    require!(
        a.shape.len() == 3 && a.shape.iter().all(|&d| d >= 0),
        "a must be a static rank-3 [B,T,H] tensor"
    );
    for idx in [1usize, 2] {
        let param = match node.input_info(idx) {
            Some(param) => param,
            None => deny!("parameter input {idx} lacks tensor type/shape info"),
        };
        require!(
            param.dtype == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT
                && param.shape == [a.shape[2]],
            "dt_bias and decay_scale must be float32 [H]"
        );
    }
    let decay = match node.output_info(0) {
        Some(decay) => decay,
        None => deny!("decay output lacks tensor type/shape info"),
    };
    require!(
        decay.dtype == a.dtype && decay.shape == a.shape,
        "decay output must match a"
    );
    if node.output_present(1) {
        require!(
            node.input_present(3),
            "b is required when beta output is requested"
        );
        let b = match node.input_info(3) {
            Some(b) => b,
            None => deny!("b lacks tensor type/shape info"),
        };
        let beta = match node.output_info(1) {
            Some(beta) => beta,
            None => deny!("beta output lacks tensor type/shape info"),
        };
        require!(
            b.dtype == a.dtype
                && b.shape == a.shape
                && beta.dtype == a.dtype
                && beta.shape == a.shape,
            "b and beta must match a"
        );
    }
    Ok(())
}

fn gated_rms_norm_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let x = ctx.resolve(&n.inputs[0])?;
    let scale = ctx.resolve(&n.inputs[1])?;
    let gate = ctx.resolve(&n.inputs[2])?;
    let out_dt = ctx.dtype_of(x);
    let x_shape = ctx.shape_of(x);
    let c = ctx.dim(scale, 0);
    let rows = ctx.size_of(x) as i32 / c;

    let x32 = ctx.astype(x, mlx::mlx_dtype__MLX_FLOAT32)?;
    let x2 = ctx.reshape(x32, &[rows, c])?;
    let squared = ctx.mul(x2, x2)?;
    let mean = ctx.emit(|res, s| unsafe { mlx::mlx_mean_axis(res, squared, 1, true, s) })?;
    let eps = ctx.scalar_f32(n.floats.get("epsilon").copied().unwrap_or(1e-5));
    let denom = ctx.add(mean, eps)?;
    let inv = ctx.unary(mlx::mlx_rsqrt, denom)?;
    let normalized = ctx.mul(x2, inv)?;
    let scale32 = ctx.astype(scale, mlx::mlx_dtype__MLX_FLOAT32)?;
    let normalized = ctx.mul(normalized, scale32)?;

    let gate32 = ctx.astype(gate, mlx::mlx_dtype__MLX_FLOAT32)?;
    let gate2 = ctx.reshape(gate32, &[rows, c])?;
    let sigmoid = ctx.unary(mlx::mlx_sigmoid, gate2)?;
    let silu = ctx.mul(gate2, sigmoid)?;
    let out32 = ctx.mul(normalized, silu)?;
    let out = ctx.astype(out32, out_dt)?;
    let out = ctx.reshape(out, &x_shape)?;
    ctx.bind(&n.outputs[0], out);
    Ok(())
}

fn gated_rms_norm_claim(node: &NodeView) -> ClaimResult {
    require!(
        node.num_inputs() == 3 && node.num_outputs() == 1,
        "GatedRMSNorm expects 3 inputs and 1 output"
    );
    let (x, scale, gate, out) = match (
        node.input_info(0),
        node.input_info(1),
        node.input_info(2),
        node.output_info(0),
    ) {
        (Some(x), Some(scale), Some(gate), Some(out)) => (x, scale, gate, out),
        _ => deny!("GatedRMSNorm requires typed input/output tensors"),
    };
    require!(
        is_mlx_float(x.dtype)
            && scale.dtype == x.dtype
            && gate.dtype == x.dtype
            && out.dtype == x.dtype,
        "X, scale, gate, and Y must share one MLX float dtype"
    );
    require!(
        !x.shape.is_empty() && x.shape.iter().all(|&d| d > 0),
        "X must have a static positive shape"
    );
    require!(
        gate.shape == x.shape && out.shape == x.shape,
        "gate and Y must have the same shape as X"
    );
    require!(
        scale.shape.len() == 1
            && scale.shape[0] > 0
            && x.shape.last().unwrap() % scale.shape[0] == 0,
        "scale must be [C] and X's last dimension must be a multiple of C"
    );
    Ok(())
}

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

pub fn register(registry: &mut OpRegistry) {
    registry.register(OpRegistration {
        domain: "",
        op_type: "TensorScatter",
        min_opset: K_ANY_OPSET,
        max_opset: K_ANY_OPSET,
        handler: tensor_scatter_op as OpHandler,
        claim: tensor_scatter_claim as ClaimPredicate,
    });
    registry.register(OpRegistration {
        domain: "com.microsoft",
        op_type: "CausalConvWithState",
        min_opset: K_ANY_OPSET,
        max_opset: K_ANY_OPSET,
        handler: causal_conv_op as OpHandler,
        claim: causal_conv_claim as ClaimPredicate,
    });
    registry.register(OpRegistration {
        domain: "",
        op_type: "CausalConvWithState",
        min_opset: 27,
        max_opset: K_ANY_OPSET,
        handler: causal_conv_op as OpHandler,
        claim: causal_conv_claim as ClaimPredicate,
    });
    registry.register(OpRegistration {
        domain: "com.microsoft",
        op_type: "LinearAttention",
        min_opset: K_ANY_OPSET,
        max_opset: K_ANY_OPSET,
        handler: linear_attention_op as OpHandler,
        claim: linear_attention_claim as ClaimPredicate,
    });
    registry.register(OpRegistration {
        domain: "",
        op_type: "LinearAttention",
        min_opset: 27,
        max_opset: K_ANY_OPSET,
        handler: linear_attention_op as OpHandler,
        claim: linear_attention_standard_claim as ClaimPredicate,
    });
    for (op_type, handler, claim) in [
        (
            "GatedAdd",
            gated_add_op as OpHandler,
            gated_add_claim as ClaimPredicate,
        ),
        (
            "LinearAttentionGate",
            linear_attention_gate_op as OpHandler,
            linear_attention_gate_claim as ClaimPredicate,
        ),
        (
            "GatedRMSNorm",
            gated_rms_norm_op as OpHandler,
            gated_rms_norm_claim as ClaimPredicate,
        ),
    ] {
        registry.register(OpRegistration {
            domain: "com.microsoft",
            op_type,
            min_opset: K_ANY_OPSET,
            max_opset: K_ANY_OPSET,
            handler,
            claim,
        });
    }
}