onnxruntime-ep-mlx 0.27.1

//! Quantized / integer op handlers — the fp32 quant path used by the cpu-recipe decoder and the
//! ai.onnx integer/QLinear coverage. Faithful port of the C++ `ops/quant.cc`, `ops/quantize.cc` and
//! `ops/quant2.cc`:
//!
//!   * MatMulNBits — int4 block-quantized weight matmul. BOTH the 3-input SYMMETRIC form
//!     (implicit zero point 8) AND the 4-input ASYMMETRIC form with an explicit packed-int4
//!     `zero_points` input (uint8). Dequant is w = (q - zp) * scale. For MLX-supported group sizes
//!     (32/64/128) the weight is repacked ONCE to MLX affine uint32 words + per-block scales/biases
//!     and run through `mlx_quantized_matmul` (weights stay compressed — the decode memory win). For
//!     block_size 16 (which `mlx_quantized_matmul` does not support) the weight is dequantized
//!     in-graph to fp32 [N,K] and a dense `mlx_matmul` runs.
//!   * GatherBlockQuantized — int4 block-quantized embedding gather + dequant, SYMMETRIC (zp=8) and
//!     ASYMMETRIC (explicit packed-int4 `zero_points`) forms.
//!   * QuantizeLinear — y = saturate(round(x / scale) + zero_point).
//!   * DequantizeLinear — y = (x - zero_point) * scale.
//!   * DynamicQuantizeLinear — compute affine uint8 scale + zero point from x's [min,0]..[max,0] range.
//!   * MatMulInteger         — (A - a_zp) @ (B - b_zp), int32 (exact via a small-K fp32 GEMM gate).
//!   * ConvInteger           — conv(x - x_zp, w - w_zp), int32 (exact via a small-accumulation gate).
//!   * QLinearMatMul         — dequant a,b -> matmul -> requantize to int8/uint8.
//!   * QLinearConv           — dequant x,w -> conv (+int32 bias) -> requantize to int8/uint8.
//!
//! Quantized weights arrive as constant initializers (surfaced as constant ctx inputs); the MatMulNBits
//! weight repack runs once and is cached on the Plan (keyed by weight name).

use std::os::raw::c_void;

use crate::engine::{MlxError, NodeDesc, Src, TranslationContext};
use crate::mlx::{Array, VectorArray};
use crate::registry::{
    ClaimResult, K_ANY_OPSET, NodeView, OpRegistration, OpRegistry, is_int_index, is_mlx_float,
};
use crate::sys::mlx;
use crate::sys::ort;
use crate::{deny, require};

// ---- small MLX helpers (each keeps + returns the raw result) -------------------------------------

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 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 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 div(
    ctx: &mut TranslationContext,
    a: mlx::mlx_array,
    b: mlx::mlx_array,
) -> Result<mlx::mlx_array, MlxError> {
    ctx.binary(mlx::mlx_divide, a, b)
}

/// round-half-to-even (matches ONNX / ORT CPU rounding).
fn round_e(ctx: &mut TranslationContext, a: mlx::mlx_array) -> Result<mlx::mlx_array, MlxError> {
    ctx.emit(|res, s| unsafe { mlx::mlx_round(res, a, 0, s) })
}

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

/// Full reduction to a scalar.
fn reduce_max(ctx: &mut TranslationContext, a: mlx::mlx_array) -> Result<mlx::mlx_array, MlxError> {
    ctx.emit(|res, s| unsafe { mlx::mlx_max(res, a, false, s) })
}
fn reduce_min(ctx: &mut TranslationContext, a: mlx::mlx_array) -> Result<mlx::mlx_array, MlxError> {
    ctx.emit(|res, s| unsafe { mlx::mlx_min(res, a, false, s) })
}

fn f32(ctx: &mut TranslationContext, v: f32) -> mlx::mlx_array {
    ctx.scalar_f32(v)
}

/// A kept 0-D uint32 scalar (bit masks / shift amounts for nibble unpacking).
fn u32_scalar(ctx: &mut TranslationContext, v: u32) -> mlx::mlx_array {
    let sh: [i32; 0] = [];
    ctx.keep(Array::from_data(
        &v as *const u32 as *const c_void,
        &sh,
        mlx::mlx_dtype__MLX_UINT32,
    ))
}

/// A node input slot is present (not an omitted optional).
fn present(n: &NodeDesc, i: usize) -> bool {
    i < n.inputs.len() && n.inputs[i].source != Src::Absent
}

/// Reshape a 1-D per-axis parameter (scale / zero point) of length L to a rank-`rank` shape that is 1
/// on every axis except `axis` (= L), so it broadcasts against the data tensor. A scalar (size <= 1)
/// parameter already broadcasts and is returned unchanged.
fn align_per_axis(
    ctx: &mut TranslationContext,
    param: mlx::mlx_array,
    rank: usize,
    axis: usize,
) -> Result<mlx::mlx_array, MlxError> {
    let len = ctx.size_of(param);
    if len <= 1 {
        return Ok(param);
    }
    let mut shape = vec![1i32; rank.max(1)];
    let ax = axis.min(shape.len() - 1);
    shape[ax] = len as i32;
    ctx.reshape(param, &shape)
}

fn norm_axis(n: &NodeDesc, rank: usize) -> usize {
    let mut axis = n.ints.get("axis").copied().unwrap_or(1);
    let r = rank as i64;
    if axis < 0 {
        axis += r;
    }
    if axis < 0 {
        axis = 0;
    }
    if axis >= r {
        axis = if r > 0 { r - 1 } else { 0 };
    }
    axis as usize
}

/// The integer range and MLX dtype for a quantized ONNX element type.
fn range_for(t: ort::ONNXTensorElementDataType) -> Option<(f32, f32, mlx::mlx_dtype)> {
    use ort::*;
    #[allow(non_upper_case_globals)]
    match t {
        ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8 => {
            Some((-128.0, 127.0, mlx::mlx_dtype__MLX_INT8))
        }
        ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8 => {
            Some((0.0, 255.0, mlx::mlx_dtype__MLX_UINT8))
        }
        ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16 => {
            Some((-32768.0, 32767.0, mlx::mlx_dtype__MLX_INT16))
        }
        ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT16 => {
            Some((0.0, 65535.0, mlx::mlx_dtype__MLX_UINT16))
        }
        _ => None,
    }
}

// ---- shared int4 unpack / block-broadcast helpers ------------------------------------------------

/// Gather rows `idx` (0-axis) of `src` → [BS, ...].
fn gather_rows(
    ctx: &mut TranslationContext,
    src: mlx::mlx_array,
    idx: mlx::mlx_array,
) -> Result<mlx::mlx_array, MlxError> {
    ctx.emit(|res, s| unsafe { mlx::mlx_take_axis(res, src, idx, 0, s) })
}

/// Unpack the interleaved low/high int4 nibbles of a packed uint8 tensor [BS, P] into the flattened
/// int4 values [BS, 2P] (uint32): column order low(byte0), high(byte0), low(byte1), high(byte1), …
/// This is the nibble layout both the packed int4 weight `data` and the packed int4 `zero_points`
/// use along the quantize axis.
fn unpack_nibbles(
    ctx: &mut TranslationContext,
    packed_u8: mlx::mlx_array,
) -> Result<mlx::mlx_array, MlxError> {
    let sh = ctx.shape_of(packed_u8);
    let bs = sh[0];
    let p = sh[1];
    let g32 = ctx.astype(packed_u8, mlx::mlx_dtype__MLX_UINT32)?;
    let mask = u32_scalar(ctx, 0x0F);
    let low = ctx.binary(mlx::mlx_bitwise_and, g32, mask)?;
    let four = u32_scalar(ctx, 4);
    let hi_sh = ctx.binary(mlx::mlx_right_shift, g32, four)?;
    let mask2 = u32_scalar(ctx, 0x0F);
    let high = ctx.binary(mlx::mlx_bitwise_and, hi_sh, mask2)?;

    let mut pair = VectorArray::new();
    pair.append(low);
    pair.append(high);
    let stacked = ctx.emit(|res, s| unsafe { mlx::mlx_stack_axis(res, pair.as_raw(), 2, s) })?; // [BS,P,2]
    drop(pair);
    ctx.reshape(stacked, &[bs, p * 2])
}

/// Broadcast a per-block float tensor [BS, nblocks] up to per-element [BS, nblocks*block]: element j
/// of a row picks its block value from block j/block.
fn broadcast_blocks(
    ctx: &mut TranslationContext,
    blocks: mlx::mlx_array,
    bs: i32,
    nblocks: i32,
    block: i32,
) -> Result<mlx::mlx_array, MlxError> {
    let r = ctx.reshape(blocks, &[bs, nblocks, 1])?;
    let bshape = [bs, nblocks, block];
    let b = ctx.emit(|res, s| unsafe {
        mlx::mlx_broadcast_to(res, r, bshape.as_ptr(), bshape.len(), s)
    })?;
    ctx.reshape(b, &[bs, nblocks * block])
}

// ---- MatMulNBits --------------------------------------------------------------------------------

/// The float dtype to run an in-graph dequant + dense matmul in: the activation's own float dtype
/// when it is fp16/bf16 (halving weight bytes + speeding the matmul), otherwise fp32.
fn compute_float_dtype(act_dt: mlx::mlx_dtype) -> mlx::mlx_dtype {
    if act_dt == mlx::mlx_dtype__MLX_FLOAT16 || act_dt == mlx::mlx_dtype__MLX_BFLOAT16 {
        act_dt
    } else {
        mlx::mlx_dtype__MLX_FLOAT32
    }
}

/// Repack our uint8 [N, nblocks, block/2] int4 weight to MLX affine uint32 words [N, K/8] (8 nibbles
/// per word, low→high along K). Cached (constant) or kept (dynamic).
fn matmulnbits_repack(
    ctx: &mut TranslationContext,
    n: &NodeDesc,
    big_n: i64,
    k: i64,
    block: i64,
) -> Result<mlx::mlx_array, MlxError> {
    let wref = &n.inputs[1];
    let key = format!("{}#qw", wref.name);
    if (wref.constant || wref.source == Src::Initializer)
        && let Some(w) = ctx.cache_get(&key)
    {
        return Ok(w);
    }
    let host = ctx.raw_host(wref)?;
    let src = host.data as *const u8;
    let nblocks = k / block;
    let blob = block / 2;
    let words = (k / 8) as usize;
    let mut packed = vec![0u32; (big_n as usize) * words];
    if !src.is_null() {
        let bytes = unsafe { std::slice::from_raw_parts(src, host.count) };
        for row in 0..big_n {
            for kk in 0..k {
                let blk = kk / block;
                let within = kk % block;
                let byte = within / 2;
                let nib = within % 2;
                let idx = ((row * nblocks + blk) * blob + byte) as usize;
                let b = bytes[idx];
                let q = if nib == 0 {
                    (b & 0x0F) as u32
                } else {
                    (b >> 4) as u32
                };
                let word = (row as usize) * words + (kk / 8) as usize;
                packed[word] |= q << (((kk % 8) as u32) * 4);
            }
        }
    }
    let sh = [big_n as i32, words as i32];
    let arr = Array::from_data(
        packed.as_ptr() as *const c_void,
        &sh,
        mlx::mlx_dtype__MLX_UINT32,
    );
    if wref.constant || wref.source == Src::Initializer {
        Ok(ctx.cache_put(key, arr))
    } else {
        Ok(ctx.keep(arr))
    }
}

/// Per-block int4 zero points [N, nblocks] as fp32 (asymmetric form): unpack the packed uint8
/// `zero_points` input and trim any trailing padding nibble.
fn matmulnbits_zp_f32(
    ctx: &mut TranslationContext,
    n: &NodeDesc,
    big_n: i32,
    nblocks: i32,
) -> Result<mlx::mlx_array, MlxError> {
    let zp_in = ctx.resolve(&n.inputs[3])?; // uint8, packed int4 (N * ceil(nblocks/2) elements)
    let cols_per_row = (nblocks + 1) / 2; // ceil(nblocks/2)
    let zp_packed = ctx.reshape(zp_in, &[big_n, cols_per_row])?; // [N, ceil(nblocks/2)]
    let zp_un = unpack_nibbles(ctx, zp_packed)?; // [N, 2*ceil(nblocks/2)]
    let cols = ctx.shape_of(zp_un)[1];
    let zp_un = if cols != nblocks {
        let start = [0i32, 0];
        let stop = [big_n, nblocks];
        let strides = [1i32, 1];
        let sl = ctx.emit(|res, s| unsafe {
            mlx::mlx_slice(
                res,
                zp_un,
                start.as_ptr(),
                start.len(),
                stop.as_ptr(),
                stop.len(),
                strides.as_ptr(),
                strides.len(),
                s,
            )
        })?;
        ctx.contiguous(sl)?
    } else {
        zp_un
    };
    ctx.astype(zp_un, mlx::mlx_dtype__MLX_FLOAT32)
}

fn matmulnbits_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let k = *n.ints.get("K").ok_or("MatMulNBits: missing K")?;
    let big_n = *n.ints.get("N").ok_or("MatMulNBits: missing N")?;
    let block = n.ints.get("block_size").copied().unwrap_or(32);
    if block <= 0 || k % block != 0 {
        return Err("MatMulNBits: bad block_size".to_string());
    }
    let nblocks = (k / block) as i32;
    let has_zp = present(n, 3);

    let a = ctx.resolve(&n.inputs[0])?;
    let ashape = ctx.shape_of(a);
    let mut m: i32 = 1;
    for &d in &ashape[..ashape.len().saturating_sub(1)] {
        m *= d;
    }
    let a2 = ctx.reshape(a, &[m, k as i32])?;
    let act_dt = ctx.dtype_of(a2);

    let supported = block == 32 || block == 64 || block == 128;
    let y = if supported {
        ctx.mark_fast("mlx_quantized_matmul");
        // Fast path: repacked uint32 weight + per-block scales/biases through mlx_quantized_matmul.
        let w = matmulnbits_repack(ctx, n, big_n, k, block)?;
        let scales = ctx.resolve(&n.inputs[2])?;
        let scales2d = ctx.reshape(scales, &[big_n as i32, nblocks])?;
        let biases = if has_zp {
            // bias = -(zp * scale)
            let zpf = matmulnbits_zp_f32(ctx, n, big_n as i32, nblocks)?;
            let neg_zp = ctx.unary(mlx::mlx_negative, zpf)?;
            mul(ctx, neg_zp, scales2d)?
        } else {
            let neg8 = f32(ctx, -8.0);
            mul(ctx, scales2d, neg8)?
        };
        let gs = mlx::mlx_optional_int_ {
            value: block as i32,
            has_value: true,
        };
        let bb = mlx::mlx_optional_int_ {
            value: 4,
            has_value: true,
        };
        let mode = c"affine".as_ptr();
        ctx.emit(|res, s| unsafe {
            mlx::mlx_quantized_matmul(res, a2, w, scales2d, biases, true, gs, bb, mode, s)
        })?
    } else {
        // Fallback (block_size mlx_quantized_matmul cannot handle, e.g. 16): dequantize in-graph.
        // Dequant + dense matmul run in the activation's float dtype (fp16/bf16 halves the weight
        // bytes + speeds the matmul; fp32 activations keep fp32 — no change).
        let comp_dt = compute_float_dtype(act_dt);
        ctx.mark_composed(format!(
            "block_size {block} unsupported by mlx_quantized_matmul → dequant + dense matmul ({})",
            crate::engine::dtype_name(comp_dt)
        ));
        let wpacked = ctx.resolve(&n.inputs[1])?; // uint8 [N, nblocks, block/2]
        let wflat = ctx.reshape(wpacked, &[big_n as i32, (k / 2) as i32])?;
        let q = unpack_nibbles(ctx, wflat)?; // uint32 [N, K]
        let qf = ctx.astype(q, comp_dt)?;
        let centered = if has_zp {
            let zpf = matmulnbits_zp_f32(ctx, n, big_n as i32, nblocks)?;
            let zpf = ctx.astype(zpf, comp_dt)?;
            let zp_full = broadcast_blocks(ctx, zpf, big_n as i32, nblocks, block as i32)?;
            sub(ctx, qf, zp_full)?
        } else {
            let eight_f = f32(ctx, 8.0);
            let eight = ctx.astype(eight_f, comp_dt)?;
            sub(ctx, qf, eight)?
        };
        let scales = ctx.resolve(&n.inputs[2])?;
        let scales2d = ctx.reshape(scales, &[big_n as i32, nblocks])?;
        let scales2d = ctx.astype(scales2d, comp_dt)?;
        let sc_full = broadcast_blocks(ctx, scales2d, big_n as i32, nblocks, block as i32)?;
        let wdeq = mul(ctx, centered, sc_full)?; // [N, K] comp_dt
        let wt = ctx.transpose(wdeq, &[1, 0])?; // [K, N]
        let a2c = ctx.astype(a2, comp_dt)?;
        ctx.binary(mlx::mlx_matmul, a2c, wt)?
    };

    // Restore leading dims with N as the last dim.
    let mut oshape: Vec<i32> = ashape.clone();
    if let Some(last) = oshape.last_mut() {
        *last = big_n as i32;
    } else {
        oshape.push(big_n as i32);
    }
    let out = ctx.reshape(y, &oshape)?;
    ctx.bind(&n.outputs[0], out);
    Ok(())
}

// ---- GatherBlockQuantized -----------------------------------------------------------------------

fn gather_block_quantized_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let block = n.ints.get("block_size").copied().unwrap_or(32) as i32;
    let idx_in = ctx.resolve(&n.inputs[1])?;
    let data = ctx.resolve(&n.inputs[0])?; // uint8 [V, D/2]
    let scales = ctx.resolve(&n.inputs[2])?; // f32 [V, nblocks]

    let ish = ctx.shape_of(idx_in);
    let mut bs: i32 = 1;
    for &d in &ish {
        bs *= d;
    }
    let idx_r = ctx.reshape(idx_in, &[bs])?;
    let idx = ctx.astype(idx_r, mlx::mlx_dtype__MLX_INT32)?;

    let g = gather_rows(ctx, data, idx)?; // [BS, D/2] uint8
    let sg = gather_rows(ctx, scales, idx)?; // [BS, nblocks]

    let packed = ctx.shape_of(g)[1];
    let d = packed * 2;
    let nblocks = d / block;

    let q = unpack_nibbles(ctx, g)?; // [BS, D] uint32
    let qf = ctx.astype(q, mlx::mlx_dtype__MLX_FLOAT32)?;

    let centered = if present(n, 3) {
        let zp_data = ctx.resolve(&n.inputs[3])?; // uint8 [V, nblocks/2] packed int4
        let zpg = gather_rows(ctx, zp_data, idx)?; // [BS, nblocks/2]
        let zp_un = unpack_nibbles(ctx, zpg)?; // [BS, 2*(nblocks/2)]
        let zp_un = if ctx.shape_of(zp_un)[1] != nblocks {
            let start = [0i32, 0];
            let stop = [bs, nblocks];
            let strides = [1i32, 1];
            let sl = ctx.emit(|res, s| unsafe {
                mlx::mlx_slice(
                    res,
                    zp_un,
                    start.as_ptr(),
                    start.len(),
                    stop.as_ptr(),
                    stop.len(),
                    strides.as_ptr(),
                    strides.len(),
                    s,
                )
            })?;
            ctx.contiguous(sl)?
        } else {
            zp_un
        };
        let zpf = ctx.astype(zp_un, mlx::mlx_dtype__MLX_FLOAT32)?;
        let zp_full = broadcast_blocks(ctx, zpf, bs, nblocks, block)?;
        sub(ctx, qf, zp_full)?
    } else {
        let eight = f32(ctx, 8.0);
        sub(ctx, qf, eight)?
    };

    let sc_full = broadcast_blocks(ctx, sg, bs, nblocks, block)?;
    let w = mul(ctx, centered, sc_full)?;

    let mut oshape: Vec<i32> = ish.clone();
    oshape.push(d);
    let out = ctx.reshape(w, &oshape)?;
    ctx.bind(&n.outputs[0], out);
    Ok(())
}

// ---- QuantizeLinear -----------------------------------------------------------------------------

fn quantize_linear_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let x0 = ctx.resolve(&n.inputs[0])?;
    let x = ctx.astype(x0, mlx::mlx_dtype__MLX_FLOAT32)?;
    let s0 = ctx.resolve(&n.inputs[1])?;
    let mut scale = ctx.astype(s0, mlx::mlx_dtype__MLX_FLOAT32)?;

    let rank = ctx.shape_of(x).len();
    let axis = norm_axis(n, rank);
    scale = align_per_axis(ctx, scale, rank, axis)?;

    let d = div(ctx, x, scale)?;
    let mut q = round_e(ctx, d)?;
    if present(n, 2) {
        let z0 = ctx.resolve(&n.inputs[2])?;
        let zp = ctx.astype(z0, mlx::mlx_dtype__MLX_FLOAT32)?;
        let zp = align_per_axis(ctx, zp, rank, axis)?;
        q = add(ctx, q, zp)?;
    }
    let (lo, hi, dt) = range_for(n.outputs[0].otype).ok_or("QuantizeLinear: bad output dtype")?;
    let flo = f32(ctx, lo);
    let fhi = f32(ctx, hi);
    let q = clip(ctx, q, flo, fhi)?;
    let out = ctx.astype(q, dt)?;
    ctx.bind(&n.outputs[0], out);
    Ok(())
}

// ---- DequantizeLinear ---------------------------------------------------------------------------

fn dequantize_linear_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let x0 = ctx.resolve(&n.inputs[0])?;
    let mut x = ctx.astype(x0, mlx::mlx_dtype__MLX_FLOAT32)?;
    let s0 = ctx.resolve(&n.inputs[1])?;
    let mut scale = ctx.astype(s0, mlx::mlx_dtype__MLX_FLOAT32)?;

    let rank = ctx.shape_of(x).len();
    let axis = norm_axis(n, rank);

    if present(n, 2) {
        let z0 = ctx.resolve(&n.inputs[2])?;
        let zp = ctx.astype(z0, mlx::mlx_dtype__MLX_FLOAT32)?;
        let zp = align_per_axis(ctx, zp, rank, axis)?;
        x = sub(ctx, x, zp)?;
    }
    scale = align_per_axis(ctx, scale, rank, axis)?;
    let out = mul(ctx, x, scale)?;
    ctx.bind(&n.outputs[0], out);
    Ok(())
}

// ---- DynamicQuantizeLinear ----------------------------------------------------------------------

fn dynamic_quantize_linear_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let x0 = ctx.resolve(&n.inputs[0])?;
    let x = ctx.astype(x0, mlx::mlx_dtype__MLX_FLOAT32)?;

    let zero = f32(ctx, 0.0);
    let rmax = reduce_max(ctx, x)?;
    let xmax = ctx.binary(mlx::mlx_maximum, rmax, zero)?;
    let rmin = reduce_min(ctx, x)?;
    let xmin = ctx.binary(mlx::mlx_minimum, rmin, zero)?;
    let span = sub(ctx, xmax, xmin)?;
    let f255 = f32(ctx, 255.0);
    let scale = div(ctx, span, f255)?;

    let lo = f32(ctx, 0.0);
    let hi = f32(ctx, 255.0);
    let zero2 = f32(ctx, 0.0);
    let neg_min = sub(ctx, zero2, xmin)?;
    let zpq = div(ctx, neg_min, scale)?;
    let zpr = round_e(ctx, zpq)?;
    let zpf = clip(ctx, zpr, lo, hi)?;

    let xq = div(ctx, x, scale)?;
    let xr = round_e(ctx, xq)?;
    let yq = add(ctx, xr, zpf)?;
    let yc = clip(ctx, yq, lo, hi)?;

    let y = ctx.astype(yc, mlx::mlx_dtype__MLX_UINT8)?;
    ctx.bind(&n.outputs[0], y);
    ctx.bind(&n.outputs[1], scale);
    let zpu = ctx.astype(zpf, mlx::mlx_dtype__MLX_UINT8)?;
    ctx.bind(&n.outputs[2], zpu);
    Ok(())
}

// ---- MatMulInteger ------------------------------------------------------------------------------

fn matmul_integer_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let a0 = ctx.resolve(&n.inputs[0])?;
    let mut a = ctx.astype(a0, mlx::mlx_dtype__MLX_FLOAT32)?;
    let b0 = ctx.resolve(&n.inputs[1])?;
    let mut b = ctx.astype(b0, mlx::mlx_dtype__MLX_FLOAT32)?;

    if present(n, 2) {
        let z0 = ctx.resolve(&n.inputs[2])?;
        let mut azp = ctx.astype(z0, mlx::mlx_dtype__MLX_FLOAT32)?;
        let sz = ctx.size_of(azp);
        if sz > 1 {
            azp = ctx.reshape(azp, &[sz as i32, 1])?;
        }
        a = sub(ctx, a, azp)?;
    }
    if present(n, 3) {
        let z0 = ctx.resolve(&n.inputs[3])?;
        let mut bzp = ctx.astype(z0, mlx::mlx_dtype__MLX_FLOAT32)?;
        let sz = ctx.size_of(bzp);
        if sz > 1 {
            bzp = ctx.reshape(bzp, &[1, sz as i32])?;
        }
        b = sub(ctx, b, bzp)?;
    }

    let y = ctx.binary(mlx::mlx_matmul, a, b)?;
    let yr = round_e(ctx, y)?;
    let out = ctx.astype(yr, mlx::mlx_dtype__MLX_INT32)?;
    ctx.bind(&n.outputs[0], out);
    Ok(())
}

// ---- conv NCHW<->NHWC transforms (mlx convs are channels-last) ----------------------------------

fn to_channels_last(
    ctx: &mut TranslationContext,
    x: mlx::mlx_array,
    spatial_rank: usize,
) -> Result<mlx::mlx_array, MlxError> {
    let t = if spatial_rank == 1 {
        ctx.transpose(x, &[0, 2, 1])?
    } else {
        ctx.transpose(x, &[0, 2, 3, 1])?
    };
    ctx.contiguous(t)
}

fn from_channels_last(
    ctx: &mut TranslationContext,
    x: mlx::mlx_array,
    spatial_rank: usize,
) -> Result<mlx::mlx_array, MlxError> {
    let t = if spatial_rank == 1 {
        ctx.transpose(x, &[0, 2, 1])?
    } else {
        ctx.transpose(x, &[0, 3, 1, 2])?
    };
    ctx.contiguous(t)
}

fn weight_to_channels_last(
    ctx: &mut TranslationContext,
    w: mlx::mlx_array,
    spatial_rank: usize,
) -> Result<mlx::mlx_array, MlxError> {
    let t = if spatial_rank == 1 {
        ctx.transpose(w, &[0, 2, 1])?
    } else {
        ctx.transpose(w, &[0, 2, 3, 1])?
    };
    ctx.contiguous(t)
}

fn attr_or(n: &NodeDesc, name: &str, size: usize, value: i64) -> Vec<i64> {
    match n.int_arrays.get(name) {
        Some(v) if !v.is_empty() => v.clone(),
        _ => vec![value; size],
    }
}

/// Centered fp32 conv of an integer x/w pair: subtract zero points, widen, transform to channels-last,
/// conv, transform back to NCHW. `x_zp` is per-tensor scalar; `w_zp` is scalar or per-output-channel.
#[allow(clippy::too_many_arguments)]
fn centered_conv(
    ctx: &mut TranslationContext,
    n: &NodeDesc,
    x: mlx::mlx_array,
    w: mlx::mlx_array,
    has_x_zp: bool,
    x_zp: mlx::mlx_array,
    has_w_zp: bool,
    w_zp: mlx::mlx_array,
    spatial_rank: usize,
) -> Result<mlx::mlx_array, MlxError> {
    let mut x = ctx.astype(x, mlx::mlx_dtype__MLX_FLOAT32)?;
    let mut w = ctx.astype(w, mlx::mlx_dtype__MLX_FLOAT32)?;
    if has_x_zp {
        let xz = ctx.astype(x_zp, mlx::mlx_dtype__MLX_FLOAT32)?;
        x = sub(ctx, x, xz)?;
    }
    if has_w_zp {
        let wrank = ctx.shape_of(w).len();
        let wz = ctx.astype(w_zp, mlx::mlx_dtype__MLX_FLOAT32)?;
        let wz = align_per_axis(ctx, wz, wrank, 0)?; // per-output-channel = weight axis 0
        w = sub(ctx, w, wz)?;
    }

    let strides = attr_or(n, "strides", spatial_rank, 1);
    let pads = attr_or(n, "pads", 2 * spatial_rank, 0);
    let dilations = attr_or(n, "dilations", spatial_rank, 1);
    let group = n.ints.get("group").copied().unwrap_or(1) as i32;

    let xl = to_channels_last(ctx, x, spatial_rank)?;
    let wl = weight_to_channels_last(ctx, w, spatial_rank)?;
    let out = if spatial_rank == 1 {
        ctx.emit(|res, s| unsafe {
            mlx::mlx_conv1d(
                res,
                xl,
                wl,
                strides[0] as i32,
                pads[0] as i32,
                dilations[0] as i32,
                group,
                s,
            )
        })?
    } else {
        ctx.emit(|res, s| unsafe {
            mlx::mlx_conv2d(
                res,
                xl,
                wl,
                strides[0] as i32,
                strides[1] as i32,
                pads[0] as i32,
                pads[1] as i32,
                dilations[0] as i32,
                dilations[1] as i32,
                group,
                s,
            )
        })?
    };
    from_channels_last(ctx, out, spatial_rank)
}

// ---- ConvInteger --------------------------------------------------------------------------------

fn conv_integer_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let x = ctx.resolve(&n.inputs[0])?;
    let w = ctx.resolve(&n.inputs[1])?;
    let spatial_rank = ctx.shape_of(x).len() - 2;

    let has_x_zp = present(n, 2);
    let has_w_zp = present(n, 3);
    let x_zp = if has_x_zp {
        ctx.resolve(&n.inputs[2])?
    } else {
        x
    };
    let w_zp = if has_w_zp {
        ctx.resolve(&n.inputs[3])?
    } else {
        w
    };

    let out = centered_conv(ctx, n, x, w, has_x_zp, x_zp, has_w_zp, w_zp, spatial_rank)?;
    let r = round_e(ctx, out)?;
    let out = ctx.astype(r, mlx::mlx_dtype__MLX_INT32)?;
    ctx.bind(&n.outputs[0], out);
    Ok(())
}

// ---- QLinearMatMul ------------------------------------------------------------------------------
//
// The dequant->matmul->requantize intermediate stays fp32 on purpose. Although the final output is
// requantized to int8/uint8 (so an fp16 intermediate looked near-free), the matmul here accumulates
// the UN-scaled centered integers: each product is O(2^8 * 2^8) and summed over K, the accumulator
// routinely exceeds the fp16 max (65504) — even at K=16 full-range int8 inputs overflow to inf,
// which then corrupts the requantized output. Measured: fp16 here is unsafe, so we keep fp32.

fn qlinear_matmul_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let a0 = ctx.resolve(&n.inputs[0])?;
    let mut a = ctx.astype(a0, mlx::mlx_dtype__MLX_FLOAT32)?;
    let asc0 = ctx.resolve(&n.inputs[1])?;
    let mut a_scale = ctx.astype(asc0, mlx::mlx_dtype__MLX_FLOAT32)?;
    let b0 = ctx.resolve(&n.inputs[3])?;
    let mut b = ctx.astype(b0, mlx::mlx_dtype__MLX_FLOAT32)?;
    let bsc0 = ctx.resolve(&n.inputs[4])?;
    let mut b_scale = ctx.astype(bsc0, mlx::mlx_dtype__MLX_FLOAT32)?;
    let ysc0 = ctx.resolve(&n.inputs[6])?;
    let mut y_scale = ctx.astype(ysc0, mlx::mlx_dtype__MLX_FLOAT32)?;

    if present(n, 2) {
        let z0 = ctx.resolve(&n.inputs[2])?;
        let mut azp = ctx.astype(z0, mlx::mlx_dtype__MLX_FLOAT32)?;
        let sz = ctx.size_of(azp);
        if sz > 1 {
            azp = ctx.reshape(azp, &[sz as i32, 1])?;
        }
        a = sub(ctx, a, azp)?;
    }
    if present(n, 5) {
        let z0 = ctx.resolve(&n.inputs[5])?;
        let mut bzp = ctx.astype(z0, mlx::mlx_dtype__MLX_FLOAT32)?;
        let sz = ctx.size_of(bzp);
        if sz > 1 {
            bzp = ctx.reshape(bzp, &[1, sz as i32])?;
        }
        b = sub(ctx, b, bzp)?;
    }

    let acc = ctx.binary(mlx::mlx_matmul, a, b)?;

    let sz = ctx.size_of(a_scale);
    if sz > 1 {
        a_scale = ctx.reshape(a_scale, &[sz as i32, 1])?;
    }
    let sz = ctx.size_of(b_scale);
    if sz > 1 {
        b_scale = ctx.reshape(b_scale, &[1, sz as i32])?;
    }
    let sz = ctx.size_of(y_scale);
    if sz > 1 {
        y_scale = ctx.reshape(y_scale, &[1, sz as i32])?;
    }

    let m1 = mul(ctx, acc, a_scale)?;
    let m2 = mul(ctx, m1, b_scale)?;
    let scaled = div(ctx, m2, y_scale)?;
    let mut q = round_e(ctx, scaled)?;
    if present(n, 7) {
        let z0 = ctx.resolve(&n.inputs[7])?;
        let mut yzp = ctx.astype(z0, mlx::mlx_dtype__MLX_FLOAT32)?;
        let sz = ctx.size_of(yzp);
        if sz > 1 {
            yzp = ctx.reshape(yzp, &[1, sz as i32])?;
        }
        q = add(ctx, q, yzp)?;
    }

    let (lo, hi, dt) = range_for(n.outputs[0].otype).ok_or("QLinearMatMul: bad output dtype")?;
    let flo = f32(ctx, lo);
    let fhi = f32(ctx, hi);
    let q = clip(ctx, q, flo, fhi)?;
    let out = ctx.astype(q, dt)?;
    ctx.bind(&n.outputs[0], out);
    Ok(())
}

// ---- QLinearConv --------------------------------------------------------------------------------
//
// Like QLinearMatMul, the conv intermediate stays fp32: the un-scaled centered-integer accumulation
// over the receptive field overflows the fp16 range (see the QLinearMatMul note), so fp16 is unsafe
// despite the int8/uint8 requantized output.

fn qlinear_conv_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let x = ctx.resolve(&n.inputs[0])?;
    let xsc0 = ctx.resolve(&n.inputs[1])?;
    let x_scale = ctx.astype(xsc0, mlx::mlx_dtype__MLX_FLOAT32)?;
    let w = ctx.resolve(&n.inputs[3])?;
    let wsc0 = ctx.resolve(&n.inputs[4])?;
    let w_scale = ctx.astype(wsc0, mlx::mlx_dtype__MLX_FLOAT32)?;
    let ysc0 = ctx.resolve(&n.inputs[6])?;
    let y_scale = ctx.astype(ysc0, mlx::mlx_dtype__MLX_FLOAT32)?;
    let spatial_rank = ctx.shape_of(x).len() - 2;

    let has_x_zp = present(n, 2);
    let has_w_zp = present(n, 5);
    let x_zp = if has_x_zp {
        ctx.resolve(&n.inputs[2])?
    } else {
        x
    };
    let w_zp = if has_w_zp {
        ctx.resolve(&n.inputs[5])?
    } else {
        w
    };

    let mut acc = centered_conv(ctx, n, x, w, has_x_zp, x_zp, has_w_zp, w_zp, spatial_rank)?;
    let out_rank = ctx.shape_of(acc).len();

    if present(n, 8) {
        let b0 = ctx.resolve(&n.inputs[8])?;
        let bias = ctx.astype(b0, mlx::mlx_dtype__MLX_FLOAT32)?;
        let bias = align_per_axis(ctx, bias, out_rank, 1)?;
        acc = add(ctx, acc, bias)?;
    }

    let xw = mul(ctx, x_scale, w_scale)?;
    let mult = div(ctx, xw, y_scale)?;
    let mult = align_per_axis(ctx, mult, out_rank, 1)?;
    let scaled = mul(ctx, acc, mult)?;
    let mut q = round_e(ctx, scaled)?;
    if present(n, 7) {
        let z0 = ctx.resolve(&n.inputs[7])?;
        let yzp = ctx.astype(z0, mlx::mlx_dtype__MLX_FLOAT32)?;
        let yzp = align_per_axis(ctx, yzp, out_rank, 1)?;
        q = add(ctx, q, yzp)?;
    }

    let (lo, hi, dt) = range_for(n.outputs[0].otype).ok_or("QLinearConv: bad output dtype")?;
    let flo = f32(ctx, lo);
    let fhi = f32(ctx, hi);
    let q = clip(ctx, q, flo, fhi)?;
    let out = ctx.astype(q, dt)?;
    ctx.bind(&n.outputs[0], out);
    Ok(())
}

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

const MAX_EXACT_ACCUM: i64 = 256;

fn is_uint8(t: ort::ONNXTensorElementDataType) -> bool {
    t == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8
}
fn is_float(t: ort::ONNXTensorElementDataType) -> bool {
    t == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT
}
fn is_int8or(t: ort::ONNXTensorElementDataType) -> bool {
    t == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8
        || t == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8
}
fn is_quant_output(t: ort::ONNXTensorElementDataType) -> bool {
    t == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8
        || t == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8
        || t == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16
        || t == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT16
}
fn is_dequant_input(t: ort::ONNXTensorElementDataType) -> bool {
    is_quant_output(t) || t == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32
}

/// A zero-point / scale parameter is claimable when absent, or present with dtype `want` and shape
/// scalar or 1-D of length `axis_len` (per-axis). `axis_len < 0` disables per-axis (scalar only).
fn param_ok(
    node: &NodeView,
    i: usize,
    want: ort::ONNXTensorElementDataType,
    axis_len: i64,
) -> bool {
    if !node.input_present(i) {
        return true;
    }
    let info = match node.input_info(i) {
        Some(x) => x,
        None => return false,
    };
    if info.dtype != want {
        return false;
    }
    if info.shape.is_empty() || (info.shape.len() == 1 && info.shape[0] == 1) {
        return true;
    }
    axis_len >= 0 && info.shape.len() == 1 && info.shape[0] == axis_len
}

fn static_positive(shape: &[i64]) -> bool {
    !shape.is_empty() && shape.iter().all(|&d| d > 0)
}

fn matmulnbits_claim(node: &NodeView) -> ClaimResult {
    let nin = node.num_inputs();
    require!(node.num_outputs() >= 1, "requires at least one output");
    let out = match node.output_info(0) {
        Some(o) => o,
        None => deny!("missing output type/shape info"),
    };
    require!(
        is_mlx_float(out.dtype),
        "output must be a float (fp32/fp16/bf16), got {}",
        crate::registry::ort_dtype_name(out.dtype)
    );
    let (a, b, s) = match (node.input_info(0), node.input_info(1), node.input_info(2)) {
        (Some(a), Some(b), Some(s)) => (a, b, s),
        _ => deny!("missing tensor type/shape info on an input"),
    };
    require!(
        is_mlx_float(a.dtype) && is_uint8(b.dtype) && is_mlx_float(s.dtype),
        "activation/scales must be float (fp32/fp16/bf16) and packed weights uint8, got a={}, b={}, scales={}",
        crate::registry::ort_dtype_name(a.dtype),
        crate::registry::ort_dtype_name(b.dtype),
        crate::registry::ort_dtype_name(s.dtype)
    );
    // The in-graph dequant + matmul run in the activation's float dtype, so activation, scales and
    // output must share one float dtype (q4f16 = all fp16; q4 = all fp32).
    require!(
        a.dtype == s.dtype && a.dtype == out.dtype,
        "activation, scales and output must share one float dtype, got a={}, scales={}, out={}",
        crate::registry::ort_dtype_name(a.dtype),
        crate::registry::ort_dtype_name(s.dtype),
        crate::registry::ort_dtype_name(out.dtype)
    );
    // Only the 3-input symmetric or 4-input asymmetric (uint8 packed int4 zero_points) forms.
    if node.input_present(3) {
        match node.input_info(3) {
            Some(zp) if is_uint8(zp.dtype) => {}
            _ => deny!("zero_points must be uint8 when present"),
        }
    }
    // Reject g_idx / bias (any present input beyond slot 3) — left to ORT CPU.
    for i in 4..nin {
        if node.input_present(i) {
            deny!("g_idx and bias inputs are not supported");
        }
    }
    let bits = node.int_attr("bits", 4);
    let block = node.int_attr("block_size", 32);
    require!(bits == 4, "only 4-bit weights are supported");
    require!(
        matches!(block, 16 | 32 | 64 | 128),
        "block_size must be 16, 32, 64, or 128 (got {block})"
    );
    Ok(())
}

fn gather_block_quantized_claim(node: &NodeView) -> ClaimResult {
    require!(node.num_outputs() >= 1, "requires at least one output");
    let out = match node.output_info(0) {
        Some(o) => o,
        None => deny!("missing output type/shape info"),
    };
    let (data, idx, scales) = match (node.input_info(0), node.input_info(1), node.input_info(2)) {
        (Some(d), Some(i), Some(s)) => (d, i, s),
        _ => deny!("missing tensor type/shape info on an input"),
    };
    if !is_uint8(data.dtype)
        || !is_int_index(idx.dtype)
        || scales.dtype != out.dtype
        || !is_mlx_float(scales.dtype)
    {
        deny!("data must be uint8, indices int32/int64, and scales/output the same float dtype");
    }
    if node.input_present(3) {
        match node.input_info(3) {
            Some(zp) if is_uint8(zp.dtype) => {}
            _ => deny!("zero_points must be uint8 when present"),
        }
    }
    require!(
        node.int_attr("bits", 4) == 4,
        "only 4-bit data is supported"
    );
    require!(
        node.int_attr("gather_axis", 0) == 0 && node.int_attr("quantize_axis", 1) == 1,
        "only gather_axis=0 and quantize_axis=1 are supported"
    );
    require!(
        node.int_attr("block_size", 128) >= 16,
        "block_size must be at least 16"
    );
    Ok(())
}

fn quantize_linear_claim(node: &NodeView) -> ClaimResult {
    let nin = node.num_inputs();
    require!(
        (2..=3).contains(&nin) && node.num_outputs() >= 1,
        "expects 2 or 3 inputs and at least one output"
    );
    let (x, s, o) = match (node.input_info(0), node.input_info(1), node.output_info(0)) {
        (Some(x), Some(s), Some(o)) => (x, s, o),
        _ => deny!("missing tensor type/shape info on an input or the output"),
    };
    if !is_float(x.dtype) || !is_float(s.dtype) {
        deny!(
            "input and scale dtypes must be float32; got {} and {}",
            crate::registry::ort_dtype_name(x.dtype),
            crate::registry::ort_dtype_name(s.dtype)
        );
    }
    if s.shape.len() > 1 || !is_quant_output(o.dtype) {
        deny!(
            "scale must be scalar or rank-1 and output dtype {} must be int8, uint8, int16, or uint16",
            crate::registry::ort_dtype_name(o.dtype)
        );
    }
    if node.input_present(2) {
        match node.input_info(2) {
            Some(z) if z.dtype == o.dtype && z.shape.len() <= 1 => {}
            _ => deny!(
                "zero_point must match output dtype {} and be scalar or rank-1",
                crate::registry::ort_dtype_name(o.dtype)
            ),
        }
    }
    Ok(())
}

fn dequantize_linear_claim(node: &NodeView) -> ClaimResult {
    let nin = node.num_inputs();
    require!(
        (2..=3).contains(&nin) && node.num_outputs() >= 1,
        "expects 2 or 3 inputs and at least one output"
    );
    let (x, s, o) = match (node.input_info(0), node.input_info(1), node.output_info(0)) {
        (Some(x), Some(s), Some(o)) => (x, s, o),
        _ => deny!("missing tensor type/shape info on an input or the output"),
    };
    if !is_dequant_input(x.dtype) || !is_float(s.dtype) || !is_float(o.dtype) {
        deny!(
            "input dtype {} must be int8, uint8, int16, uint16, or int32; scale/output must be float32 (got {} -> {})",
            crate::registry::ort_dtype_name(x.dtype),
            crate::registry::ort_dtype_name(s.dtype),
            crate::registry::ort_dtype_name(o.dtype)
        );
    }
    if s.shape.len() > 1 {
        deny!("scale must be scalar or rank-1");
    }
    if node.input_present(2) {
        match node.input_info(2) {
            Some(z) if z.dtype == x.dtype && z.shape.len() <= 1 => {}
            _ => deny!(
                "zero_point must match input dtype {} and be scalar or rank-1",
                crate::registry::ort_dtype_name(x.dtype)
            ),
        }
    }
    Ok(())
}

fn dynamic_quantize_linear_claim(node: &NodeView) -> ClaimResult {
    require!(
        node.num_inputs() == 1 && node.num_outputs() == 3,
        "expects 1 input and 3 outputs, got {}in/{}out",
        node.num_inputs(),
        node.num_outputs()
    );
    let (x, y, sc, z) = match (
        node.input_info(0),
        node.output_info(0),
        node.output_info(1),
        node.output_info(2),
    ) {
        (Some(x), Some(y), Some(sc), Some(z)) => (x, y, sc, z),
        _ => deny!("missing tensor type/shape info on an input or output"),
    };
    require!(
        is_float(x.dtype) && is_uint8(y.dtype) && is_float(sc.dtype) && is_uint8(z.dtype),
        "input and scale must be float32; quantized output and zero point must be uint8"
    );
    Ok(())
}

fn matmul_integer_claim(node: &NodeView) -> ClaimResult {
    let nin = node.num_inputs();
    require!(
        (2..=4).contains(&nin) && node.num_outputs() >= 1,
        "expects 2 to 4 inputs and at least one output"
    );
    let (a, b, o) = match (node.input_info(0), node.input_info(1), node.output_info(0)) {
        (Some(a), Some(b), Some(o)) => (a, b, o),
        _ => deny!("missing tensor type/shape info on an input or the output"),
    };
    if !is_int8or(a.dtype)
        || !is_int8or(b.dtype)
        || o.dtype != ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32
    {
        deny!(
            "input dtypes {} and {} must be int8 or uint8, and output dtype {} must be int32",
            crate::registry::ort_dtype_name(a.dtype),
            crate::registry::ort_dtype_name(b.dtype),
            crate::registry::ort_dtype_name(o.dtype)
        );
    }
    if a.shape.len() != 2 || b.shape.len() != 2 {
        deny!("inputs must both be rank-2 matrices");
    }
    let k = a.shape[1];
    if k <= 0 || k != b.shape[0] || k > MAX_EXACT_ACCUM {
        deny!("inner dimensions must match and K must be 1..={MAX_EXACT_ACCUM}");
    }
    if node.input_present(2) {
        match node.input_info(2) {
            Some(z) if z.dtype == a.dtype && z.shape.len() <= 1 => {}
            _ => deny!(
                "a_zero_point must match A dtype {} and be scalar or rank-1",
                crate::registry::ort_dtype_name(a.dtype)
            ),
        }
    }
    if node.input_present(3) {
        match node.input_info(3) {
            Some(z) if z.dtype == b.dtype && z.shape.len() <= 1 => {}
            _ => deny!(
                "b_zero_point must match B dtype {} and be scalar or rank-1",
                crate::registry::ort_dtype_name(b.dtype)
            ),
        }
    }
    Ok(())
}

fn conv_attrs_ok(
    node: &NodeView,
    spatial_rank: usize,
    w_shape: &[i64],
    channels: i64,
    group: i64,
) -> bool {
    if node.string_attr("auto_pad", "NOTSET") != "NOTSET" {
        return false;
    }
    if group <= 0 || channels % group != 0 || w_shape[1] != channels / group {
        return false;
    }
    if w_shape[0] % group != 0 {
        return false;
    }
    let get = |name: &str, def_len: usize, def: i64| -> Option<Vec<i64>> {
        let (present, v) = node.ints_attr(name);
        if !present {
            Some(vec![def; def_len])
        } else if v.len() == def_len {
            Some(v)
        } else {
            None
        }
    };
    let strides = match get("strides", spatial_rank, 1) {
        Some(v) => v,
        None => return false,
    };
    let dilations = match get("dilations", spatial_rank, 1) {
        Some(v) => v,
        None => return false,
    };
    let pads = match get("pads", 2 * spatial_rank, 0) {
        Some(v) => v,
        None => return false,
    };
    if strides.iter().any(|&v| v <= 0)
        || dilations.iter().any(|&v| v <= 0)
        || pads.iter().any(|&v| v < 0)
    {
        return false;
    }
    for i in 0..spatial_rank {
        if pads[i] != pads[i + spatial_rank] {
            return false; // mlx conv takes a symmetric pad only
        }
    }
    let (kpresent, kernel) = node.ints_attr("kernel_shape");
    if kpresent {
        if kernel.len() != spatial_rank {
            return false;
        }
        for i in 0..spatial_rank {
            if kernel[i] != w_shape[i + 2] {
                return false;
            }
        }
    }
    true
}

fn conv_accum_exact(w_shape: &[i64]) -> bool {
    let mut n_acc = w_shape[1];
    for &d in &w_shape[2..] {
        n_acc *= d;
    }
    (1..=MAX_EXACT_ACCUM).contains(&n_acc)
}

fn conv_integer_claim(node: &NodeView) -> ClaimResult {
    let nin = node.num_inputs();
    require!(
        (2..=4).contains(&nin) && node.num_outputs() == 1,
        "expects 2 to 4 inputs and 1 output"
    );
    let (x, w, o) = match (node.input_info(0), node.input_info(1), node.output_info(0)) {
        (Some(x), Some(w), Some(o)) => (x, w, o),
        _ => deny!("missing tensor type/shape info on an input or the output"),
    };
    if !is_int8or(x.dtype)
        || !is_int8or(w.dtype)
        || o.dtype != ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32
    {
        deny!(
            "input dtypes {} and {} must be int8 or uint8, and output dtype {} must be int32",
            crate::registry::ort_dtype_name(x.dtype),
            crate::registry::ort_dtype_name(w.dtype),
            crate::registry::ort_dtype_name(o.dtype)
        );
    }
    if (x.shape.len() != 3 && x.shape.len() != 4) || w.shape.len() != x.shape.len() {
        deny!("input/weight must have matching rank, with input rank 3 or 4");
    }
    if !static_positive(&x.shape) || !static_positive(&w.shape) {
        deny!("input and weight shapes must be static and positive");
    }
    let spatial_rank = x.shape.len() - 2;
    let group = node.int_attr("group", 1);
    require!(
        conv_attrs_ok(node, spatial_rank, &w.shape, x.shape[1], group),
        "convolution attributes, channels, group, or kernel shape are unsupported"
    );
    require!(
        conv_accum_exact(&w.shape),
        "convolution accumulation size must be 1..={MAX_EXACT_ACCUM}"
    );
    require!(
        param_ok(node, 2, x.dtype, -1),
        "x_zero_point must match input dtype {} and be scalar",
        crate::registry::ort_dtype_name(x.dtype)
    );
    require!(
        param_ok(node, 3, w.dtype, w.shape[0]),
        "w_zero_point must match weight dtype {} and be scalar or per-output-channel",
        crate::registry::ort_dtype_name(w.dtype)
    );
    Ok(())
}

fn qlinear_matmul_claim(node: &NodeView) -> ClaimResult {
    require!(
        node.num_inputs() == 8 && node.num_outputs() == 1,
        "expects 8 inputs and 1 output, got {}in/{}out",
        node.num_inputs(),
        node.num_outputs()
    );
    let (a, b, o) = match (node.input_info(0), node.input_info(3), node.output_info(0)) {
        (Some(a), Some(b), Some(o)) => (a, b, o),
        _ => deny!("missing tensor type/shape info on an input or the output"),
    };
    if !is_int8or(a.dtype) || !is_int8or(b.dtype) || !is_int8or(o.dtype) {
        deny!("A, B, and output must each be int8 or uint8");
    }
    if a.shape.len() != 2 || b.shape.len() != 2 {
        deny!("A and B must be rank-2 matrices");
    }
    let (m, k, big_n) = (a.shape[0], a.shape[1], b.shape[1]);
    if k <= 0 || k != b.shape[0] || k > MAX_EXACT_ACCUM || m <= 0 || big_n <= 0 {
        deny!(
            "matrix dimensions must be positive, inner dimensions match, and K must be 1..={MAX_EXACT_ACCUM}"
        );
    }
    let f = ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT;
    require!(
        param_ok(node, 1, f, m),
        "a_scale must be float32 and scalar or length M"
    );
    require!(
        param_ok(node, 2, a.dtype, m),
        "a_zero_point must match A dtype and be scalar or length M"
    );
    require!(
        param_ok(node, 4, f, big_n),
        "b_scale must be float32 and scalar or length N"
    );
    require!(
        param_ok(node, 5, b.dtype, big_n),
        "b_zero_point must match B dtype and be scalar or length N"
    );
    require!(
        param_ok(node, 6, f, big_n),
        "y_scale must be float32 and scalar or length N"
    );
    require!(
        param_ok(node, 7, o.dtype, big_n),
        "y_zero_point must match output dtype and be scalar or length N"
    );
    Ok(())
}

fn qlinear_conv_claim(node: &NodeView) -> ClaimResult {
    let nin = node.num_inputs();
    require!(
        (8..=9).contains(&nin) && node.num_outputs() == 1,
        "expects 8 or 9 inputs and 1 output"
    );
    let (x, w, o) = match (node.input_info(0), node.input_info(3), node.output_info(0)) {
        (Some(x), Some(w), Some(o)) => (x, w, o),
        _ => deny!("missing tensor type/shape info on an input or the output"),
    };
    if !is_int8or(x.dtype) || !is_int8or(w.dtype) || !is_int8or(o.dtype) {
        deny!("input, weight, and output must each be int8 or uint8");
    }
    if (x.shape.len() != 3 && x.shape.len() != 4) || w.shape.len() != x.shape.len() {
        deny!("input/weight must have matching rank, with input rank 3 or 4");
    }
    if !static_positive(&x.shape) || !static_positive(&w.shape) {
        deny!("input and weight shapes must be static and positive");
    }
    let spatial_rank = x.shape.len() - 2;
    let big_m = w.shape[0];
    let group = node.int_attr("group", 1);
    require!(
        conv_attrs_ok(node, spatial_rank, &w.shape, x.shape[1], group),
        "convolution attributes, channels, group, or kernel shape are unsupported"
    );
    require!(
        conv_accum_exact(&w.shape),
        "convolution accumulation size must be 1..={MAX_EXACT_ACCUM}"
    );
    let f = ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT;
    require!(param_ok(node, 1, f, -1), "x_scale must be a scalar float32");
    require!(
        param_ok(node, 2, x.dtype, -1),
        "x_zero_point must match input dtype and be scalar"
    );
    require!(
        param_ok(node, 4, f, big_m),
        "w_scale must be float32 and scalar or per-output-channel"
    );
    require!(
        param_ok(node, 5, w.dtype, big_m),
        "w_zero_point must match weight dtype and be scalar or per-output-channel"
    );
    require!(param_ok(node, 6, f, -1), "y_scale must be a scalar float32");
    require!(
        param_ok(node, 7, o.dtype, -1),
        "y_zero_point must match output dtype and be scalar"
    );
    if node.input_present(8) {
        match node.input_info(8) {
            Some(bi)
                if bi.dtype
                    == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32
                    && bi.shape.len() == 1
                    && bi.shape[0] == big_m => {}
            _ => deny!("bias must be int32 with one element per output channel"),
        }
    }
    Ok(())
}

// ---- QMoE ---------------------------------------------------------------------------------------
// Quantized Mixture-of-Experts (com.microsoft.QMoE). Supported subset: quant_type='int', symmetric
// int4/int8 column- or block-wise expert weights, top-k softmax routing, and SwiGLU (interleaved,
// swiglu_fusion=1) / silu / gelu / relu / identity activation. Compute is DENSE over experts: every
// expert is evaluated for every token and then masked by the top-k routing gate (weights that are
// zero outside a token's top-k). That avoids any data-dependent token scatter/gather, so the op is a
// handful of batched MLX matmuls that fuse into the compiled closure. Declined (left to ORT CPU):
// fc3 experts, asymmetric zero_points, router_weights, sparse-mixer, and fp4/fp8/wfp4afp8 quant.

const QMOE_IN_INPUT: usize = 0;
const QMOE_IN_ROUTER: usize = 1;
const QMOE_IN_FC1_W: usize = 2;
const QMOE_IN_FC1_S: usize = 3;
const QMOE_IN_FC1_B: usize = 4;
const QMOE_IN_FC2_W: usize = 5;
const QMOE_IN_FC2_S: usize = 6;
const QMOE_IN_FC2_B: usize = 7;
const QMOE_IN_FC3_W: usize = 8;
const QMOE_IN_FC1_ZP: usize = 11;
const QMOE_IN_FC2_ZP: usize = 12;
const QMOE_IN_ROUTER_W: usize = 14;

/// Dequantize a symmetric int4/int8 expert-weight tensor [E, N, ceil(K/pack)] to a dense [E, N, K]
/// MLX array in `comp_dt`. Faithful to ORT's `DequantizeBlock`: value = (q - 2^(bits-1)) * scale,
/// where the nibble/byte for column c is at packed index c/pack with shift (c%pack)*bits, and each
/// scale spans `block_size` columns along K (column-wise = one scale per row when block_size==0).
#[allow(clippy::too_many_arguments)] // full expert-weight geometry (E/N/K/bits/block) + ctx/indices
fn qmoe_dequant(
    ctx: &mut TranslationContext,
    n: &NodeDesc,
    widx: usize,
    sidx: usize,
    e: i64,
    nrows: i64,
    k: i64,
    bits: i64,
    block_size: i64,
    comp_dt: mlx::mlx_dtype,
) -> Result<mlx::mlx_array, MlxError> {
    let pack = 8 / bits;
    if block_size > 0 && k % block_size != 0 {
        return Err(format!(
            "QMoE: K={k} not divisible by block_size={block_size}"
        ));
    }
    let blocks = if block_size > 0 { k / block_size } else { 1 };
    let block = if block_size > 0 { block_size } else { k };
    let rows = (e * nrows) as i32;
    let packed_k = (k + pack - 1) / pack;
    let wp = ctx.resolve(&n.inputs[widx])?;
    let w2d = ctx.reshape(wp, &[rows, packed_k as i32])?;
    let q = if bits == 4 {
        let un = unpack_nibbles(ctx, w2d)?; // [rows, 2*ceil(K/2)]
        let cols = ctx.shape_of(un)[1];
        if cols != k as i32 {
            ctx.slice(un, &[0, 0], &[rows, k as i32])?
        } else {
            un
        }
    } else {
        w2d // 8-bit: one byte per value, already [rows, K]
    };
    let qf = ctx.astype(q, comp_dt)?;
    let zpv = f32(ctx, (1i64 << (bits - 1)) as f32);
    let zp = ctx.astype(zpv, comp_dt)?;
    let centered = sub(ctx, qf, zp)?;
    let scales = ctx.resolve(&n.inputs[sidx])?;
    let sc2d = ctx.reshape(scales, &[rows, blocks as i32])?;
    let sc2d = ctx.astype(sc2d, comp_dt)?;
    let sc_full = broadcast_blocks(ctx, sc2d, rows, blocks as i32, block as i32)?;
    let deq = mul(ctx, centered, sc_full)?;
    ctx.reshape(deq, &[e as i32, nrows as i32, k as i32])
}

/// SwiGLU on interleaved gate/linear halves (matches ORT `ApplySwiGLUActivation`):
/// g = min(gate, limit); l = clamp(linear, -limit, limit); out = g*sigmoid(alpha*g) * (l + beta).
/// The clamp (limit=+inf), the alpha scale (alpha=1) and the beta shift (beta=0) are identities in
/// the standard SwiGLU form, so we skip those dispatches — fewer eager kernels, same result.
fn qmoe_swiglu(
    ctx: &mut TranslationContext,
    gate: mlx::mlx_array,
    linear: mlx::mlx_array,
    alpha: f32,
    beta: f32,
    limit: f32,
    comp_dt: mlx::mlx_dtype,
) -> Result<mlx::mlx_array, MlxError> {
    let (g, l) = if limit.is_finite() {
        let lim_v = f32(ctx, limit);
        let lim = ctx.astype(lim_v, comp_dt)?;
        let neglim_v = f32(ctx, -limit);
        let neglim = ctx.astype(neglim_v, comp_dt)?;
        let g = ctx.binary(mlx::mlx_minimum, gate, lim)?;
        let l = ctx.emit(|res, s| unsafe { mlx::mlx_clip(res, linear, neglim, lim, s) })?;
        (g, l)
    } else {
        (gate, linear)
    };
    let ag = if alpha != 1.0 {
        let alpha_v = f32(ctx, alpha);
        let alpha_c = ctx.astype(alpha_v, comp_dt)?;
        mul(ctx, g, alpha_c)?
    } else {
        g
    };
    let sig = ctx.unary(mlx::mlx_sigmoid, ag)?;
    let swish = mul(ctx, g, sig)?;
    let lb = if beta != 0.0 {
        let beta_v = f32(ctx, beta);
        let beta_c = ctx.astype(beta_v, comp_dt)?;
        add(ctx, l, beta_c)?
    } else {
        l
    };
    mul(ctx, swish, lb)
}

/// Elementwise activation for the non-SwiGLU forms (matches ORT `ApplyActivation`).
fn qmoe_activation(
    ctx: &mut TranslationContext,
    x: mlx::mlx_array,
    act: &str,
    comp_dt: mlx::mlx_dtype,
) -> Result<mlx::mlx_array, MlxError> {
    match act {
        "relu" => {
            let z_v = f32(ctx, 0.0);
            let z = ctx.astype(z_v, comp_dt)?;
            ctx.binary(mlx::mlx_maximum, x, z)
        }
        "silu" => {
            let s = ctx.unary(mlx::mlx_sigmoid, x)?;
            mul(ctx, x, s)
        }
        "gelu" => {
            // 0.5 * x * (1 + tanh(0.7978845608 * (x + 0.044715 * x^3)))
            let c0_v = f32(ctx, 0.044715);
            let c0 = ctx.astype(c0_v, comp_dt)?;
            let x2 = mul(ctx, x, x)?;
            let x3 = mul(ctx, x2, x)?;
            let cx3 = mul(ctx, x3, c0)?;
            let inner = add(ctx, x, cx3)?;
            let s_v = f32(ctx, 0.797_884_6);
            let s = ctx.astype(s_v, comp_dt)?;
            let scaled = mul(ctx, inner, s)?;
            let th = ctx.unary(mlx::mlx_tanh, scaled)?;
            let one_v = f32(ctx, 1.0);
            let one = ctx.astype(one_v, comp_dt)?;
            let onep = add(ctx, th, one)?;
            let half_v = f32(ctx, 0.5);
            let half = ctx.astype(half_v, comp_dt)?;
            let hx = mul(ctx, x, half)?;
            mul(ctx, hx, onep)
        }
        "identity" => Ok(x),
        other => Err(format!("QMoE: unsupported activation {other}")),
    }
}

fn qmoe_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let k_top = *n.ints.get("k").unwrap_or(&1) as i32;
    let bits = *n.ints.get("expert_weight_bits").unwrap_or(&4);
    let block_size = *n.ints.get("block_size").unwrap_or(&0);
    let fusion = *n.ints.get("swiglu_fusion").unwrap_or(&0);
    let act = n
        .strings
        .get("activation_type")
        .cloned()
        .unwrap_or_else(|| "relu".to_string());
    let alpha = *n.floats.get("activation_alpha").unwrap_or(&1.0);
    let beta = *n.floats.get("activation_beta").unwrap_or(&0.0);
    let limit = n
        .floats
        .get("swiglu_limit")
        .copied()
        .unwrap_or(f32::INFINITY);

    let x = ctx.resolve(&n.inputs[QMOE_IN_INPUT])?;
    let out_dt = ctx.dtype_of(x);
    let comp_dt = compute_float_dtype(out_dt);
    let xs = ctx.shape_of(x);
    if xs.is_empty() {
        return Err("QMoE: input must have rank >= 1".to_string());
    }
    let hidden = *xs.last().unwrap() as i64;
    let mut t: i64 = 1;
    for &d in &xs[..xs.len() - 1] {
        t *= d as i64;
    }

    // Expert dimensions from the (static) fc1 weight shape [E, F*I, H/pack].
    let w1_arr = ctx.resolve(&n.inputs[QMOE_IN_FC1_W])?;
    let w1s = ctx.shape_of(w1_arr);
    if w1s.len() != 3 {
        return Err("QMoE: fc1_experts_weights must be rank 3".to_string());
    }
    let e = w1s[0] as i64;
    let fi = w1s[1] as i64;
    let f = if fusion > 0 { 2 } else { 1 };
    if fi % f != 0 {
        return Err("QMoE: fc1 out features not divisible by fusion size".to_string());
    }
    let inter = fi / f;

    // Routing gate [T, E]: softmax over the top-k logits (== masked softmax with non-top-k -> -inf).
    // Threshold = the k-th largest logit per row (ascending sort, column E-k).
    let router = ctx.resolve(&n.inputs[QMOE_IN_ROUTER])?;
    let router_f = ctx.astype(router, mlx::mlx_dtype__MLX_FLOAT32)?;
    let sorted = ctx.emit(|res, s| unsafe { mlx::mlx_sort_axis(res, router_f, 1, s) })?; // [T, E] asc
    let kcol = (e as i32) - k_top;
    let thr = ctx.slice(sorted, &[0, kcol], &[t as i32, kcol + 1])?; // [T, 1] k-th largest
    let mask = ctx.binary(mlx::mlx_greater_equal, router_f, thr)?; // [T, E] bool
    let neg_inf = f32(ctx, -3.0e38);
    let masked = ctx.where_(mask, router_f, neg_inf)?;
    let gate_f = ctx.softmax_axis(masked, 1)?; // [T, E] f32
    let gate = ctx.astype(gate_f, comp_dt)?;

    // Dequantize experts to dense [E, F*I, H] and [E, H, I].
    let w1 = qmoe_dequant(
        ctx,
        n,
        QMOE_IN_FC1_W,
        QMOE_IN_FC1_S,
        e,
        fi,
        hidden,
        bits,
        block_size,
        comp_dt,
    )?;
    let w2 = qmoe_dequant(
        ctx,
        n,
        QMOE_IN_FC2_W,
        QMOE_IN_FC2_S,
        e,
        hidden,
        inter,
        bits,
        block_size,
        comp_dt,
    )?;

    // FC1: C1[e,t,:] = X[t,:] @ W1[e]^T  ->  [E, T, F*I].
    let xc = ctx.astype(x, comp_dt)?;
    let x2 = ctx.reshape(xc, &[t as i32, hidden as i32])?;
    let x_e = ctx.expand_dims(x2, 0)?; // [1, T, H]
    let bshape = [e as i32, t as i32, hidden as i32];
    let x_b = ctx.emit(|res, s| unsafe {
        mlx::mlx_broadcast_to(res, x_e, bshape.as_ptr(), bshape.len(), s)
    })?;
    let x_b = ctx.contiguous(x_b)?;
    let w1t = ctx.transpose(w1, &[0, 2, 1])?; // [E, H, F*I]
    let mut c1 = ctx.matmul(x_b, w1t)?; // [E, T, F*I]
    if present(n, QMOE_IN_FC1_B) {
        let b1 = ctx.resolve(&n.inputs[QMOE_IN_FC1_B])?;
        let b1 = ctx.astype(b1, comp_dt)?;
        let b1 = ctx.reshape(b1, &[e as i32, 1, fi as i32])?;
        c1 = add(ctx, c1, b1)?;
    }

    // Activation -> A2 [E, T, I].
    let a2 = if act == "swiglu" {
        let c1r = ctx.reshape(c1, &[e as i32, t as i32, inter as i32, 2])?;
        let gate_p = ctx.slice(c1r, &[0, 0, 0, 0], &[e as i32, t as i32, inter as i32, 1])?;
        let lin_p = ctx.slice(c1r, &[0, 0, 0, 1], &[e as i32, t as i32, inter as i32, 2])?;
        let gate_p = ctx.reshape(gate_p, &[e as i32, t as i32, inter as i32])?;
        let lin_p = ctx.reshape(lin_p, &[e as i32, t as i32, inter as i32])?;
        qmoe_swiglu(ctx, gate_p, lin_p, alpha, beta, limit, comp_dt)?
    } else {
        qmoe_activation(ctx, c1, &act, comp_dt)?
    };

    // FC2: C2[e,t,:] = A2[e,t,:] @ W2[e]^T  ->  [E, T, H].
    let w2t = ctx.transpose(w2, &[0, 2, 1])?; // [E, I, H]
    let mut c2 = ctx.matmul(a2, w2t)?; // [E, T, H]
    if present(n, QMOE_IN_FC2_B) {
        let b2 = ctx.resolve(&n.inputs[QMOE_IN_FC2_B])?;
        let b2 = ctx.astype(b2, comp_dt)?;
        let b2 = ctx.reshape(b2, &[e as i32, 1, hidden as i32])?;
        c2 = add(ctx, c2, b2)?;
    }

    // Weighted sum over experts: out[t,:] = sum_e gate[t,e] * C2[e,t,:].
    let gate_t = ctx.transpose(gate, &[1, 0])?; // [E, T]
    let gate_e = ctx.reshape(gate_t, &[e as i32, t as i32, 1])?;
    let weighted = mul(ctx, c2, gate_e)?; // [E, T, H]
    let out_e = ctx.emit(|res, s| unsafe { mlx::mlx_sum_axis(res, weighted, 0, false, s) })?; // [T, H]
    let out = ctx.astype(out_e, out_dt)?;
    let out = ctx.reshape(out, &xs)?;
    ctx.bind(&n.outputs[0], out);
    Ok(())
}

fn qmoe_claim(node: &NodeView) -> ClaimResult {
    let ninputs = node.num_inputs();
    require!(
        ninputs >= 7,
        "QMoE expects at least 7 inputs (input, router, fc1_w, fc1_s, fc1_b, fc2_w, fc2_s), got {}",
        ninputs
    );
    let out_dt = match node.output_info(0) {
        Some(o) if is_mlx_float(o.dtype) => o.dtype,
        Some(o) => deny!(
            "output must have an MLX float dtype, got {}",
            crate::registry::ort_dtype_name(o.dtype)
        ),
        None => deny!("output lacks tensor type/shape info"),
    };
    match node.input_info(QMOE_IN_INPUT) {
        Some(i) => require!(
            i.dtype == out_dt,
            "input dtype must match output dtype {}",
            crate::registry::ort_dtype_name(out_dt)
        ),
        None => deny!("input lacks tensor type/shape info"),
    }
    let qt = node.string_attr("quant_type", "int");
    require!(
        qt == "int",
        "only quant_type='int' is supported (got {:?})",
        qt
    );
    let bits = node.int_attr("expert_weight_bits", 4);
    require!(
        bits == 4 || bits == 8,
        "expert_weight_bits must be 4 or 8 (got {bits})"
    );
    require!(
        node.input_present(QMOE_IN_FC1_W) && node.input_present(QMOE_IN_FC2_W),
        "requires fc1/fc2 expert weights"
    );
    require!(
        node.input_present(QMOE_IN_FC1_S) && node.input_present(QMOE_IN_FC2_S),
        "int quant requires fc1/fc2 scales"
    );
    require!(
        !node.input_present(QMOE_IN_FC3_W),
        "fc3 experts (non-fused gated MLP) not supported"
    );
    require!(
        !node.input_present(QMOE_IN_FC1_ZP) && !node.input_present(QMOE_IN_FC2_ZP),
        "asymmetric quantization (zero_points) not supported"
    );
    require!(
        !node.input_present(QMOE_IN_ROUTER_W),
        "router_weights (decoupled select/aggregate routing) not supported"
    );
    require!(
        node.int_attr("use_sparse_mixer", 0) != 1,
        "sparse mixer routing not supported"
    );
    let act = node.string_attr("activation_type", "relu");
    require!(
        matches!(
            act.as_str(),
            "swiglu" | "silu" | "gelu" | "relu" | "identity"
        ),
        "unsupported activation_type {:?}",
        act
    );
    let fusion = node.int_attr("swiglu_fusion", 0);
    if act == "swiglu" {
        require!(
            fusion == 1,
            "SwiGLU requires swiglu_fusion=1 (interleaved), got {fusion}"
        );
    } else {
        require!(
            fusion == 0,
            "swiglu_fusion is only valid with activation_type=swiglu"
        );
    }
    let k_top = node.int_attr("k", 1);
    require!(k_top >= 1, "k must be >= 1 (got {k_top})");
    // fc1 weights must be a static rank-3 [E, F*I, H/pack] tensor, and k <= num_experts.
    match node.input_info(QMOE_IN_FC1_W) {
        Some(w) => {
            require!(
                w.shape.len() == 3 && w.shape.iter().all(|&d| d > 0),
                "fc1_experts_weights must be a static rank-3 tensor, got shape {:?}",
                w.shape
            );
            require!(
                k_top <= w.shape[0],
                "k={k_top} exceeds num_experts={}",
                w.shape[0]
            );
        }
        None => deny!("fc1_experts_weights lacks tensor type/shape info"),
    }
    Ok(())
}

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

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

pub fn register(registry: &mut OpRegistry) {
    reg(
        registry,
        "com.microsoft",
        "MatMulNBits",
        matmulnbits_op,
        matmulnbits_claim,
    );
    reg(registry, "com.microsoft", "QMoE", qmoe_op, qmoe_claim);
    reg(
        registry,
        "com.microsoft",
        "GatherBlockQuantized",
        gather_block_quantized_op,
        gather_block_quantized_claim,
    );
    reg(
        registry,
        "",
        "QuantizeLinear",
        quantize_linear_op,
        quantize_linear_claim,
    );
    reg(
        registry,
        "",
        "DequantizeLinear",
        dequantize_linear_op,
        dequantize_linear_claim,
    );
    reg(
        registry,
        "",
        "DynamicQuantizeLinear",
        dynamic_quantize_linear_op,
        dynamic_quantize_linear_claim,
    );
    reg(
        registry,
        "",
        "MatMulInteger",
        matmul_integer_op,
        matmul_integer_claim,
    );
    reg(
        registry,
        "",
        "ConvInteger",
        conv_integer_op,
        conv_integer_claim,
    );
    reg(
        registry,
        "",
        "QLinearMatMul",
        qlinear_matmul_op,
        qlinear_matmul_claim,
    );
    reg(
        registry,
        "",
        "QLinearConv",
        qlinear_conv_op,
        qlinear_conv_claim,
    );
}