onnxruntime-ep-mlx 0.31.1

//! Shape / data-movement op handlers. Faithful port of the C++ `ops/shape.cc` (+ `WhereOp` from
//! `ops/math.cc`): Gather, GatherElements, ScatterElements, Concat, Reshape, Transpose, Unsqueeze,
//! Squeeze, Flatten, Expand, Slice, Split (multi-output), Tile, Pad, Identity, Range, Shape, Size,
//! ConstantOfShape and Where.
//!
//! Shape/axes/indices operands (Reshape shape, Slice starts/ends/axes/steps, Pad pads, Expand/Tile
//! shape, ConstantOfShape/Range scalars) are read at translate time via `RawHost`/`read_ints`.
//! Live boundary parameters use guarded general traces; unguarded interior parameters stay eager.
//! View-op boundary outputs (transpose/slice/expand/split/gather) are forced
//! contiguous before the shared CopyOut memcpy. Zero-size results (Pad/Expand) are re-materialised as
//! clean zeros arrays rather than rejected to CPU.

use half::{bf16, f16};

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

// ---- small helpers -----------------------------------------------------------------------------

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

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

fn clamp_i64(v: i64, lo: i64, hi: i64) -> i64 {
    v.max(lo).min(hi)
}

fn int_attr(n: &NodeDesc, name: &str, default: i64) -> i64 {
    n.ints.get(name).copied().unwrap_or(default)
}

/// Read an int32/int64 parameter input (shape/axes/starts/...) at translate time.
fn read_ints(ctx: &mut TranslationContext, n: &NodeDesc, i: usize) -> Result<Vec<i64>, MlxError> {
    ctx.read_ints_eval(&n.inputs[i])
}

/// Read an axes/split-style list from either the opset-13 input or the opset<13 INTS attribute.
fn read_list_input_or_attr(
    ctx: &mut TranslationContext,
    n: &NodeDesc,
    input_idx: usize,
    attr: &str,
) -> Result<(bool, Vec<i64>), MlxError> {
    if present(n, input_idx) {
        return Ok((true, read_ints(ctx, n, input_idx)?));
    }
    if let Some(v) = n.int_arrays.get(attr) {
        return Ok((true, v.clone()));
    }
    Ok((false, Vec::new()))
}

fn where_op(
    ctx: &mut TranslationContext,
    cond: mlx::mlx_array,
    x: mlx::mlx_array,
    y: mlx::mlx_array,
) -> Result<mlx::mlx_array, MlxError> {
    ctx.emit(|res, s| unsafe { mlx::mlx_where(res, cond, x, y, s) })
}

fn resize_src_coord(
    mode: &str,
    output_index: i64,
    scale: f64,
    input_len: i64,
    output_len: i64,
) -> f64 {
    match mode {
        "align_corners" => {
            if output_len == 1 {
                0.0
            } else {
                output_index as f64 * (input_len - 1) as f64 / (output_len - 1) as f64
            }
        }
        "asymmetric" => output_index as f64 / scale,
        "pytorch_half_pixel" if output_len == 1 => 0.0,
        _ => (output_index as f64 + 0.5) / scale - 0.5,
    }
}

fn resize_nearest_index(mode: &str, source: f64, input_len: i64) -> i32 {
    let index = match mode {
        "floor" => source.floor(),
        "ceil" => source.ceil(),
        "round_prefer_ceil" => (source + 0.5).floor(),
        _ => (source - 0.5).ceil(),
    } as i64;
    clamp_i64(index, 0, input_len - 1) as i32
}

/// Normalize + wrap negative gather indices into [0, dim) as int32 (the take/gather index dtype).
fn normalize_indices(
    ctx: &mut TranslationContext,
    indices: mlx::mlx_array,
    dim: i32,
) -> Result<mlx::mlx_array, MlxError> {
    let idx = ctx.astype(indices, mlx::mlx_dtype__MLX_INT32)?;
    let dim_s = ctx.scalar_i32(dim);
    let zero_s = ctx.scalar_i32(0);
    let neg = ctx.binary(mlx::mlx_less, idx, zero_s)?;
    let wrapped = ctx.binary(mlx::mlx_add, idx, dim_s)?;
    where_op(ctx, neg, wrapped, idx)
}

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

fn gather_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let data = ctx.resolve(&n.inputs[0])?;
    let indices = ctx.resolve(&n.inputs[1])?;
    let rank = ctx.ndim(data) as i32;
    let axis = norm_axis(int_attr(n, "axis", 0), rank);
    let dim = ctx.dim(data, axis);
    let idx = normalize_indices(ctx, indices, dim)?;
    let r = ctx.emit(|res, s| unsafe { mlx::mlx_take_axis(res, data, idx, axis, s) })?;
    let r = ctx.contiguous(r)?;
    ctx.bind(&n.outputs[0], r);
    Ok(())
}

fn gather_elements_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let data = ctx.resolve(&n.inputs[0])?;
    let indices = ctx.resolve(&n.inputs[1])?;
    let rank = ctx.ndim(data) as i32;
    let axis = norm_axis(int_attr(n, "axis", 0), rank);
    let dim = ctx.dim(data, axis);
    let idx = normalize_indices(ctx, indices, dim)?;
    let r = ctx.emit(|res, s| unsafe { mlx::mlx_take_along_axis(res, data, idx, axis, s) })?;
    let r = ctx.contiguous(r)?;
    ctx.bind(&n.outputs[0], r);
    Ok(())
}

// ---- GatherND --------------------------------------------------------------------------------
fn i64_row(ctx: &mut TranslationContext, data: &[i64]) -> mlx::mlx_array {
    ctx.keep(Array::from_data(
        data.as_ptr() as *const std::os::raw::c_void,
        &[1, data.len() as i32],
        mlx::mlx_dtype__MLX_INT64,
    ))
}

fn gathernd_extent_product(shape: &[i32]) -> Result<i32, MlxError> {
    shape.iter().try_fold(1i32, |product, &extent| {
        product
            .checked_mul(extent)
            .ok_or_else(|| "GatherND: flattened extent exceeds MLX int32 indexing".to_string())
    })
}

fn gathernd_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let data = ctx.resolve(&n.inputs[0])?;
    let indices = ctx.resolve(&n.inputs[1])?;
    let dshape = ctx.shape_of(data);
    let ishape = ctx.shape_of(indices);
    let (r, q) = (dshape.len(), ishape.len());
    let batch_dims = usize::try_from(int_attr(n, "batch_dims", 0))
        .map_err(|_| "GatherND: batch_dims must be non-negative".to_string())?;
    if batch_dims >= r.min(q) || dshape[..batch_dims] != ishape[..batch_dims] {
        return Err("GatherND: invalid or mismatched batch dimensions".to_string());
    }
    let m = *ishape
        .last()
        .ok_or("GatherND: indices must have rank >= 1")? as usize;
    if m == 0 || m > r - batch_dims {
        return Err(format!(
            "GatherND: indices last dim {m} out of range for data rank {r} and batch_dims {batch_dims}"
        ));
    }
    let mut oshape = ishape[..q - 1].to_vec();
    oshape.extend_from_slice(&dshape[batch_dims + m..]);
    if oshape.contains(&0) {
        let out = ctx.zeros(&oshape, ctx.dtype_of(data))?;
        ctx.bind(&n.outputs[0], out);
        return Ok(());
    }
    let rows = gathernd_extent_product(&dshape[..batch_dims + m])?;
    let tail = gathernd_extent_product(&dshape[batch_dims + m..])?;
    let data_flat = ctx.reshape(data, &[rows, tail])?;
    let p = gathernd_extent_product(&ishape[..q - 1])?;
    let idx2d = ctx.reshape(indices, &[p, m as i32])?;
    let idx64 = ctx.astype(idx2d, mlx::mlx_dtype__MLX_INT64)?;

    let dims: Vec<i64> = dshape[batch_dims..batch_dims + m]
        .iter()
        .map(|&d| d as i64)
        .collect();
    let mut strides = vec![0i64; m];
    let mut acc = 1i64;
    for j in (0..m).rev() {
        strides[j] = acc;
        acc *= dims[j];
    }
    let dims_a = i64_row(ctx, &dims); // [1,m]
    let strides_a = i64_row(ctx, &strides); // [1,m]

    // Wrap negative indices per column: idx = idx < 0 ? idx + dim : idx.
    let zero = ctx.scalar_i64(0);
    let neg = ctx.binary(mlx::mlx_less, idx64, zero)?;
    let wrapped = ctx.binary(mlx::mlx_add, idx64, dims_a)?;
    let idxpos = ctx.where_(neg, wrapped, idx64)?;
    // flat row index per tuple, then gather.
    let scaled = ctx.binary(mlx::mlx_multiply, idxpos, strides_a)?; // [P,m]
    let mut flat = ctx.emit(|res, s| unsafe { mlx::mlx_sum_axis(res, scaled, 1, false, s) })?;
    if batch_dims > 0 {
        let batches = gathernd_extent_product(&dshape[..batch_dims])?;
        let per_batch = gathernd_extent_product(&ishape[batch_dims..q - 1])?;
        let batch = ctx.arange(0.0, batches as f64, 1.0, mlx::mlx_dtype__MLX_INT64)?;
        let stride = ctx.scalar_i64(acc);
        let offsets = ctx.binary(mlx::mlx_multiply, batch, stride)?;
        let offsets = ctx.reshape(offsets, &[batches, 1])?;
        let offsets = ctx.emit(|res, s| unsafe {
            mlx::mlx_broadcast_to(res, offsets, [batches, per_batch].as_ptr(), 2, s)
        })?;
        let offsets = ctx.reshape(offsets, &[p])?;
        flat = ctx.binary(mlx::mlx_add, flat, offsets)?;
    }
    let flat = ctx.astype(flat, mlx::mlx_dtype__MLX_INT32)?;
    let gathered = ctx.emit(|res, s| unsafe { mlx::mlx_take_axis(res, data_flat, flat, 0, s) })?; // [P,tail]

    let out = ctx.reshape(gathered, &oshape)?;
    let out = ctx.contiguous(out)?;
    ctx.bind(&n.outputs[0], out);
    Ok(())
}

fn gathernd_claim(node: &NodeView) -> ClaimResult {
    require!(
        node.num_inputs() == 2 && node.num_outputs() == 1,
        "GatherND expects 2 inputs and 1 output, got {}in/{}out",
        node.num_inputs(),
        node.num_outputs()
    );
    let (data, out) = match movable_io(node) {
        Some(v) => v,
        None => deny!("missing tensor type/shape info on the data input or output"),
    };
    require!(
        is_movable(data),
        "data dtype {} is not a movable dtype supported on GPU",
        crate::registry::ort_dtype_name(data)
    );
    require!(out == data, "output dtype must match data dtype");
    let idt = match node.input_info(1) {
        Some(i) => i.dtype,
        None => deny!("missing indices tensor type/shape info"),
    };
    require!(
        idt == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64
            || idt == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32,
        "GatherND indices must be int32/int64, got {}",
        crate::registry::ort_dtype_name(idt)
    );
    let data = node.input_info(0).expect("validated data above");
    let indices = node.input_info(1).expect("validated indices above");
    let batch_dims = node.int_attr("batch_dims", 0);
    require!(
        batch_dims >= 0 && batch_dims < data.shape.len().min(indices.shape.len()) as i64,
        "GatherND: batch_dims must be in [0, min(data rank, indices rank))"
    );
    let batch_dims = batch_dims as usize;
    let width = *indices.shape.last().expect("validated indices rank above");
    require!(
        width > 0 && width <= (data.shape.len() - batch_dims) as i64,
        "GatherND: tuple width must be known and in [1, data rank - batch_dims]"
    );
    for axis in 0..batch_dims {
        let (data_dim, index_dim) = (data.shape[axis], indices.shape[axis]);
        require!(
            data_dim < 0 || index_dim < 0 || data_dim == index_dim,
            "GatherND: data and indices batch dimension {axis} must match"
        );
    }
    Ok(())
}

fn scatter_elements_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let data = ctx.resolve(&n.inputs[0])?;
    let indices = ctx.resolve(&n.inputs[1])?;
    let updates = ctx.resolve(&n.inputs[2])?;
    let shape = ctx.shape_of(data);
    let index_shape = ctx.shape_of(indices);
    let update_shape = ctx.shape_of(updates);
    let rank = shape.len() as i64;
    let axis = int_attr(n, "axis", 0);
    if rank == 0 || axis < -rank || axis >= rank {
        return Err("ScatterElements: axis is out of range".to_string());
    }
    if index_shape != update_shape || index_shape.len() != shape.len() {
        return Err(
            "ScatterElements: indices/updates must have equal shapes and data rank".to_string(),
        );
    }
    let axis = norm_axis(axis, rank as i32);
    let mut extent = shape.clone();
    for i in 0..shape.len() {
        if i != axis as usize {
            if index_shape[i] > shape[i] {
                return Err("ScatterElements: non-axis indices extent exceeds data".to_string());
            }
            extent[i] = index_shape[i];
        }
    }
    if index_shape.contains(&0) {
        let result = ctx.contiguous(data)?;
        ctx.bind(&n.outputs[0], result);
        return Ok(());
    }
    let start = vec![0; shape.len()];
    let partial = extent != shape;
    let source = if partial {
        ctx.slice(data, &start, &extent)?
    } else {
        data
    };
    let idx = normalize_indices(ctx, indices, shape[axis as usize])?;
    let result = ctx.emit(|res, stream| unsafe {
        mlx::mlx_put_along_axis(res, source, idx, updates, axis, stream)
    })?;
    let result = if partial {
        ctx.slice_update(data, result, &start, &extent)?
    } else {
        result
    };
    let result = ctx.contiguous(result)?;
    ctx.bind(&n.outputs[0], result);
    Ok(())
}

fn compatible_shape(left: &[i64], right: &[i64]) -> bool {
    left.len() == right.len()
        && left
            .iter()
            .zip(right)
            .all(|(&a, &b)| a < 0 || b < 0 || a == b)
}

fn center_crop_pad_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let data = ctx.resolve(&n.inputs[0])?;
    let input_shape = ctx.shape_of(data);
    let rank = input_shape.len() as i32;
    let target = read_ints(ctx, n, 1)?;
    let axes: Vec<i64> = n
        .int_arrays
        .get("axes")
        .cloned()
        .unwrap_or_else(|| (0..rank as i64).collect());

    let mut cropped_shape = input_shape.clone();
    let mut start = vec![0i32; rank as usize];
    let mut stop = input_shape.clone();
    let mut low = vec![0i32; axes.len()];
    let mut high = vec![0i32; axes.len()];
    let mut norm_axes = vec![0i32; axes.len()];
    for (i, (&axis, &wanted)) in axes.iter().zip(&target).enumerate() {
        let ax = norm_axis(axis, rank) as usize;
        let wanted = dim_i32(wanted)?;
        norm_axes[i] = ax as i32;
        if wanted < input_shape[ax] {
            start[ax] = (input_shape[ax] - wanted) / 2;
            stop[ax] = start[ax] + wanted;
            cropped_shape[ax] = wanted;
        } else {
            low[i] = (wanted - input_shape[ax]) / 2;
            high[i] = wanted - input_shape[ax] - low[i];
        }
    }

    let cropped = if cropped_shape != input_shape {
        ctx.slice(data, &start, &stop)?
    } else {
        data
    };
    let zero_i64 = ctx.scalar_i64(0);
    let zero = ctx.astype(zero_i64, ctx.dtype_of(data))?;
    let out = ctx.emit(|res, s| unsafe {
        mlx::mlx_pad(
            res,
            cropped,
            norm_axes.as_ptr(),
            norm_axes.len(),
            low.as_ptr(),
            low.len(),
            high.as_ptr(),
            high.len(),
            zero,
            c"constant".as_ptr(),
            s,
        )
    })?;
    ctx.bind(&n.outputs[0], out);
    Ok(())
}

fn compress_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let mut data = ctx.resolve(&n.inputs[0])?;
    let mut shape = ctx.shape_of(data);
    let axis = n.ints.get("axis").copied();
    let axis = if let Some(axis) = axis {
        norm_axis(axis, shape.len() as i32)
    } else {
        data = ctx.reshape(data, &[shape.iter().product()])?;
        shape = ctx.shape_of(data);
        0
    };
    let condition = ctx.raw_host(&n.inputs[1])?;
    let values =
        unsafe { std::slice::from_raw_parts(condition.data as *const u8, condition.count) };
    let extent = shape[axis as usize] as usize;
    let indices: Vec<i32> = values
        .iter()
        .take(extent)
        .enumerate()
        .filter_map(|(i, &value)| (value != 0).then_some(i as i32))
        .collect();
    let index = ctx.keep(Array::from_data(
        indices.as_ptr() as *const std::os::raw::c_void,
        &[indices.len() as i32],
        mlx::mlx_dtype__MLX_INT32,
    ));
    let out = ctx.emit(|res, s| unsafe { mlx::mlx_take_axis(res, data, index, axis, s) })?;
    let out = ctx.contiguous(out)?;
    ctx.bind(&n.outputs[0], out);
    Ok(())
}

fn depth_to_space_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let data = ctx.resolve(&n.inputs[0])?;
    let shape = ctx.shape_of(data);
    let block = int_attr(n, "blocksize", 0) as i32;
    let channels = shape[1] / (block * block);
    let reshaped = if n.strings.get("mode").map(String::as_str).unwrap_or("DCR") == "CRD" {
        let x = ctx.reshape(
            data,
            &[shape[0], channels, block, block, shape[2], shape[3]],
        )?;
        ctx.transpose(x, &[0, 1, 4, 2, 5, 3])?
    } else {
        let x = ctx.reshape(
            data,
            &[shape[0], block, block, channels, shape[2], shape[3]],
        )?;
        ctx.transpose(x, &[0, 3, 4, 1, 5, 2])?
    };
    let out = ctx.reshape(
        reshaped,
        &[shape[0], channels, shape[2] * block, shape[3] * block],
    )?;
    let out = ctx.contiguous(out)?;
    ctx.bind(&n.outputs[0], out);
    Ok(())
}

fn eye_like_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let input = ctx.resolve(&n.inputs[0])?;
    let shape = ctx.shape_of(input);
    let dtype = mlx_dtype_from_onnx(n.outputs[0].otype);
    let k = int_attr(n, "k", 0) as i32;
    let out = ctx.emit(|res, s| unsafe { mlx::mlx_eye(res, shape[0], shape[1], k, dtype, s) })?;
    ctx.bind(&n.outputs[0], out);
    Ok(())
}

fn reverse_sequence_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let data = ctx.resolve(&n.inputs[0])?;
    let lengths = ctx.resolve(&n.inputs[1])?;
    let rank = ctx.ndim(data) as i32;
    let batch_axis = norm_axis(int_attr(n, "batch_axis", 1), rank);
    let time_axis = norm_axis(int_attr(n, "time_axis", 0), rank);

    let mut perm = vec![batch_axis, time_axis];
    perm.extend((0..rank).filter(|&axis| axis != batch_axis && axis != time_axis));
    let mut inverse = vec![0i32; rank as usize];
    for (i, &axis) in perm.iter().enumerate() {
        inverse[axis as usize] = i as i32;
    }
    let transposed = ctx.transpose(data, &perm)?;
    let shape = ctx.shape_of(transposed);
    let lens = ctx.astype(lengths, mlx::mlx_dtype__MLX_INT32)?;
    let lens = ctx.reshape(lens, &[shape[0], 1])?;
    let time = ctx.arange(0.0, shape[1] as f64, 1.0, mlx::mlx_dtype__MLX_INT32)?;
    let time = ctx.reshape(time, &[1, shape[1]])?;
    let before = ctx.binary(mlx::mlx_less, time, lens)?;
    let one = ctx.scalar_i32(1);
    let last = ctx.binary(mlx::mlx_subtract, lens, one)?;
    let reversed = ctx.binary(mlx::mlx_subtract, last, time)?;
    let indices = ctx.where_(before, reversed, time)?;
    let mut index_shape = vec![shape[0], shape[1]];
    index_shape.resize(shape.len(), 1);
    let indices = ctx.reshape(indices, &index_shape)?;
    let indices = ctx.emit(|res, s| unsafe {
        mlx::mlx_broadcast_to(res, indices, shape.as_ptr(), shape.len(), s)
    })?;
    let out =
        ctx.emit(|res, s| unsafe { mlx::mlx_take_along_axis(res, transposed, indices, 1, s) })?;
    let out = ctx.transpose(out, &inverse)?;
    let out = ctx.contiguous(out)?;
    ctx.bind(&n.outputs[0], out);
    Ok(())
}

fn scatter_nd_geometry(
    data: &[i32],
    indices: &[i32],
    updates: &[i32],
) -> Result<(usize, i32, i32, i32), MlxError> {
    let width = *indices
        .last()
        .ok_or("ScatterND: indices rank must be positive")?;
    if data.is_empty() || width <= 0 || width as usize > data.len() {
        return Err("ScatterND: index tuple width is outside the data rank".to_string());
    }
    let width = width as usize;
    let mut expected = indices[..indices.len() - 1].to_vec();
    expected.extend_from_slice(&data[width..]);
    if updates != expected || data.iter().chain(indices).any(|&d| d < 0) {
        return Err("ScatterND: updates shape does not match index/data geometry".to_string());
    }
    let product = |shape: &[i32]| {
        shape.iter().try_fold(1i32, |a, &b| {
            a.checked_mul(b)
                .ok_or_else(|| "ScatterND: flattened extent exceeds int32 indexing".to_string())
        })
    };
    Ok((
        width,
        product(&data[..width])?,
        product(&data[width..])?,
        product(&indices[..indices.len() - 1])?,
    ))
}

fn scatter_nd_add_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let data = ctx.resolve(&n.inputs[0])?;
    let indices = ctx.resolve(&n.inputs[1])?;
    let updates = ctx.resolve(&n.inputs[2])?;
    let shape = ctx.shape_of(data);
    let index_shape = ctx.shape_of(indices);
    let update_shape = ctx.shape_of(updates);
    let (width, rows, tail, count) = scatter_nd_geometry(&shape, &index_shape, &update_shape)?;
    if count == 0 {
        let out = ctx.contiguous(data)?;
        ctx.bind(&n.outputs[0], out);
        return Ok(());
    }
    // MLX scatter does not check bounds. Validate live tuples before a native write.
    // This form is eager so checks run again when values change at identical shapes.
    let raw = ctx.read_ints_eval(&n.inputs[1])?;
    if raw.len() != count as usize * width {
        return Err("ScatterND: index value count does not match its shape".to_string());
    }
    let mut flat = Vec::with_capacity(count as usize);
    for tuple in raw.chunks_exact(width) {
        let mut index = 0i64;
        for (&value, &extent) in tuple.iter().zip(&shape[..width]) {
            let extent = i64::from(extent);
            if value < -extent || value >= extent {
                return Err("ScatterND: index is out of bounds".to_string());
            }
            index = index * extent + if value < 0 { value + extent } else { value };
        }
        flat.push(index as i32);
    }
    if tail == 0 {
        let out = ctx.contiguous(data)?;
        ctx.bind(&n.outputs[0], out);
        return Ok(());
    }
    let indices = ctx.keep(Array::from_data(
        flat.as_ptr() as *const std::os::raw::c_void,
        &[count],
        mlx::mlx_dtype__MLX_INT32,
    ));
    let base = ctx.reshape(data, &[rows, tail])?;
    let updates = ctx.reshape(updates, &[count, 1, tail])?;
    let out = ctx.emit(|res, stream| unsafe {
        mlx::mlx_scatter_add_single(res, base, indices, updates, 0, stream)
    })?;
    let out = ctx.reshape(out, &shape)?;
    let out = ctx.contiguous(out)?;
    ctx.bind(&n.outputs[0], out);
    Ok(())
}

fn scatter_nd_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    if n.strings.get("reduction").map(String::as_str) == Some("add") {
        return scatter_nd_add_op(ctx, n);
    }
    let data = ctx.resolve(&n.inputs[0])?;
    let indices = ctx.resolve(&n.inputs[1])?;
    let updates = ctx.resolve(&n.inputs[2])?;
    let dshape = ctx.shape_of(data);
    let ishape = ctx.shape_of(indices);
    let ushape = ctx.shape_of(updates);
    let (m, rows, tail, p) = scatter_nd_geometry(&dshape, &ishape, &ushape)?;
    if p == 0 || tail == 0 {
        let out = ctx.contiguous(data)?;
        ctx.bind(&n.outputs[0], out);
        return Ok(());
    }
    let data_flat = ctx.reshape(data, &[rows, tail])?;
    let idx2d = ctx.reshape(indices, &[p, m as i32])?;
    let idx64 = ctx.astype(idx2d, mlx::mlx_dtype__MLX_INT64)?;

    let dims: Vec<i64> = dshape[..m].iter().map(|&d| d as i64).collect();
    let mut strides = vec![0i64; m];
    let mut acc = 1i64;
    for j in (0..m).rev() {
        strides[j] = acc;
        acc *= dims[j];
    }
    let dims_a = i64_row(ctx, &dims);
    let strides_a = i64_row(ctx, &strides);
    let zero = ctx.scalar_i64(0);
    let negative = ctx.binary(mlx::mlx_less, idx64, zero)?;
    let wrapped = ctx.binary(mlx::mlx_add, idx64, dims_a)?;
    let normalized = ctx.where_(negative, wrapped, idx64)?;
    let scaled = ctx.binary(mlx::mlx_multiply, normalized, strides_a)?;
    let flat = ctx.emit(|res, s| unsafe { mlx::mlx_sum_axis(res, scaled, 1, false, s) })?;
    let flat = ctx.astype(flat, mlx::mlx_dtype__MLX_INT32)?;
    let flat = ctx.reshape(flat, &[p, 1])?;
    let flat =
        ctx.emit(|res, s| unsafe { mlx::mlx_broadcast_to(res, flat, [p, tail].as_ptr(), 2, s) })?;
    let updates = ctx.reshape(updates, &[p, tail])?;
    let out =
        ctx.emit(|res, s| unsafe { mlx::mlx_put_along_axis(res, data_flat, flat, updates, 0, s) })?;
    let out = ctx.reshape(out, &dshape)?;
    let out = ctx.contiguous(out)?;
    ctx.bind(&n.outputs[0], out);
    Ok(())
}

fn space_to_depth_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let data = ctx.resolve(&n.inputs[0])?;
    let shape = ctx.shape_of(data);
    let block = int_attr(n, "blocksize", 0) as i32;
    let h = shape[2] / block;
    let w = shape[3] / block;
    let reshaped = ctx.reshape(data, &[shape[0], shape[1], h, block, w, block])?;
    let perm = match n.strings.get("mode").map(String::as_str).unwrap_or("DCR") {
        "DCR" => [0, 3, 5, 1, 2, 4],
        "CRD" => [0, 1, 3, 5, 2, 4],
        mode => {
            return Err(format!(
                "SpaceToDepth: mode must be DCR or CRD, got {mode:?}"
            ));
        }
    };
    let transposed = ctx.transpose(reshaped, &perm)?;
    let out = ctx.reshape(transposed, &[shape[0], shape[1] * block * block, h, w])?;
    let out = ctx.contiguous(out)?;
    ctx.bind(&n.outputs[0], out);
    Ok(())
}

fn concat_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let mut vec = VectorArray::new();
    let mut rank = 0i32;
    for (i, inp) in n.inputs.iter().enumerate() {
        let a = ctx.resolve(inp)?;
        if i == 0 {
            rank = ctx.ndim(a) as i32;
        }
        vec.append(a);
    }
    let axis = norm_axis(int_attr(n, "axis", 0), rank);
    let r = ctx.emit(|res, s| unsafe { mlx::mlx_concatenate_axis(res, vec.as_raw(), axis, s) })?;
    ctx.bind(&n.outputs[0], r);
    Ok(())
}

fn reshape_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let data = ctx.resolve(&n.inputs[0])?;
    let shape = read_ints(ctx, n, 1)?;
    let in_shape = ctx.shape_of(data);
    let allowzero = int_attr(n, "allowzero", 0) != 0;
    let target: Vec<i32> = shape
        .iter()
        .enumerate()
        .map(|(i, &d)| {
            if d == 0 && !allowzero && i < in_shape.len() {
                Ok(in_shape[i]) // allowzero=0: copy the input dim
            } else {
                dim_i32(d) // preserves -1 (infer); errors on >i32 dims
            }
        })
        .collect::<Result<_, _>>()?;
    let r = ctx.reshape(data, &target)?;
    ctx.bind(&n.outputs[0], r);
    Ok(())
}

fn checked_transpose_perm(rank: usize, perm: Option<&[i64]>) -> Result<Vec<i32>, MlxError> {
    let rank = i32::try_from(rank).map_err(|_| "Transpose rank exceeds int32".to_string())?;
    let Some(perm) = perm else {
        return Ok((0..rank).rev().collect());
    };
    if perm.len() != rank as usize {
        return Err("Transpose permutation length must equal input rank".to_string());
    }
    let mut seen = vec![false; perm.len()];
    let mut result = Vec::with_capacity(perm.len());
    for &axis in perm {
        if axis < 0 || axis >= i64::from(rank) {
            return Err(format!(
                "Transpose axis {axis} is out of range for rank {rank}"
            ));
        }
        let axis = axis as i32;
        if seen[axis as usize] {
            return Err("Transpose permutation must contain each axis exactly once".to_string());
        }
        seen[axis as usize] = true;
        result.push(axis);
    }
    Ok(result)
}

fn transpose_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let shape = ctx.tensor_shape(&n.inputs[0])?;
    if shape
        .iter()
        .any(|&dim| dim < 0 || dim > i64::from(i32::MAX))
    {
        return Err("Transpose input dimensions must fit nonnegative int32".to_string());
    }
    let perm = checked_transpose_perm(shape.len(), n.int_arrays.get("perm").map(Vec::as_slice))?;
    let data = ctx.resolve(&n.inputs[0])?;
    let mut t = ctx.transpose(data, &perm)?;
    // A very wide frozen projection becomes a tall RHS in backward. MLX GEMM
    // prefers that reduction dimension row-major; pack it once, including BF16
    // round-trip views proven immutable by the translation context.
    if !ctx.is_capturing()
        && shape.len() == 2
        && perm == [1, 0]
        && shape[0] >= 64
        && shape[1] / shape[0] >= 8
        && ctx.size_of(data) >= 1024 * 1024
        && ctx.dtype_of(data) == mlx::mlx_dtype__MLX_FLOAT32
        && std::env::var("ONNXRUNTIME_EP_MLX_PACK_FROZEN_VIEWS").as_deref() != Ok("0")
        && let Some(origin) = ctx.frozen_key(&n.inputs[0])
    {
        t = ctx.cache_contiguous(format!("__frozen_transpose__{origin}:{shape:?}"), t)?;
        ctx.mark_fast("mlx_cached_frozen_transpose");
    }
    // Other transposes remain zero-copy strided views; materialised to contiguous only if it reaches the output boundary.
    ctx.bind(&n.outputs[0], t);
    Ok(())
}

fn checked_rank_change_axes(
    shape: &[i64],
    axes: Option<&[i64]>,
    inserting: bool,
    allow_negative: bool,
) -> Result<Vec<i32>, MlxError> {
    let rank = shape
        .len()
        .checked_add(if inserting {
            axes.map_or(0, <[i64]>::len)
        } else {
            0
        })
        .and_then(|rank| i32::try_from(rank).ok())
        .ok_or_else(|| "Squeeze/Unsqueeze rank exceeds int32".to_string())?;
    let axes = match axes {
        Some(axes) => axes,
        None if inserting => return Err("Unsqueeze requires axes".to_string()),
        None => {
            return shape
                .iter()
                .enumerate()
                .filter(|&(_, &dim)| dim == 1)
                .map(|(axis, _)| {
                    i32::try_from(axis).map_err(|_| "Squeeze rank exceeds int32".to_string())
                })
                .collect();
        }
    };
    let mut normalized = Vec::with_capacity(axes.len());
    for &axis in axes {
        if (!allow_negative && axis < 0) || axis < -i64::from(rank) || axis >= i64::from(rank) {
            return Err(format!(
                "Squeeze/Unsqueeze axis {axis} is out of range for rank {rank}"
            ));
        }
        let axis = norm_axis(axis, rank);
        if !inserting && shape[axis as usize] >= 0 && shape[axis as usize] != 1 {
            return Err(format!("Squeeze axis {axis} must have extent 1"));
        }
        normalized.push(axis);
    }
    normalized.sort_unstable();
    if inserting && normalized.windows(2).any(|pair| pair[0] == pair[1]) {
        return Err("Unsqueeze axes must be unique".to_string());
    }
    normalized.dedup();
    Ok(normalized)
}

fn rank_change_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let shape = ctx.tensor_shape(&n.inputs[0])?;
    if shape
        .iter()
        .any(|&dim| dim < 0 || dim > i64::from(i32::MAX))
    {
        return Err("Squeeze/Unsqueeze dimensions must fit nonnegative int32".to_string());
    }
    if present(n, 1) && ctx.tensor_shape(&n.inputs[1])?.len() != 1 {
        return Err("Squeeze/Unsqueeze axes must be one-dimensional".to_string());
    }
    let (have_axes, axes) = read_list_input_or_attr(ctx, n, 1, "axes")?;
    let inserting = n.op_type == "Unsqueeze";
    let axes = checked_rank_change_axes(
        &shape,
        have_axes.then_some(axes.as_slice()),
        inserting,
        n.since_version >= 11,
    )?;
    let data = ctx.resolve(&n.inputs[0])?;
    let r = if axes.is_empty() {
        data
    } else if inserting {
        ctx.emit(|res, s| unsafe {
            mlx::mlx_expand_dims_axes(res, data, axes.as_ptr(), axes.len(), s)
        })?
    } else {
        ctx.emit(|res, s| unsafe {
            mlx::mlx_squeeze_axes(res, data, axes.as_ptr(), axes.len(), s)
        })?
    };
    ctx.bind(&n.outputs[0], r);
    Ok(())
}

fn checked_flatten_axis(rank: usize, axis: i64, allow_negative: bool) -> Result<usize, MlxError> {
    let rank =
        i64::from(i32::try_from(rank).map_err(|_| "Flatten rank exceeds int32".to_string())?);
    if (!allow_negative && axis < 0) || axis < -rank || axis > rank {
        return Err(format!(
            "Flatten axis {axis} is out of range for rank {rank}"
        ));
    }
    Ok((if axis < 0 { axis + rank } else { axis }) as usize)
}

fn checked_flatten_product(dims: &[i64]) -> Result<i32, MlxError> {
    if dims.iter().any(|&dim| dim < 0 || dim > i64::from(i32::MAX)) {
        return Err("Flatten input dimensions must fit nonnegative int32".to_string());
    }
    if dims.contains(&0) {
        return Ok(0);
    }
    dims.iter().try_fold(1i32, |product, &dim| {
        product
            .checked_mul(dim as i32)
            .ok_or_else(|| "Flatten folded dimension exceeds int32".to_string())
    })
}

fn checked_flatten_shape(
    shape: &[i64],
    axis: i64,
    allow_negative: bool,
) -> Result<[i32; 2], MlxError> {
    let axis = checked_flatten_axis(shape.len(), axis, allow_negative)?;
    Ok([
        checked_flatten_product(&shape[..axis])?,
        checked_flatten_product(&shape[axis..])?,
    ])
}

fn flatten_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let shape = ctx.tensor_shape(&n.inputs[0])?;
    let target = checked_flatten_shape(&shape, int_attr(n, "axis", 1), n.since_version >= 11)?;
    let data = ctx.resolve(&n.inputs[0])?;
    let r = ctx.reshape(data, &target)?;
    ctx.bind(&n.outputs[0], r);
    Ok(())
}

fn expand_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let data = ctx.resolve(&n.inputs[0])?;
    let target = read_ints(ctx, n, 1)?;
    let in_shape = ctx.shape_of(data);
    let out_rank = in_shape.len().max(target.len());
    let mut result = vec![0i32; out_rank];
    for i in 0..out_rank {
        let d_in = if i + in_shape.len() < out_rank {
            1i64
        } else {
            in_shape[i - (out_rank - in_shape.len())] as i64
        };
        let d_t = if i + target.len() < out_rank {
            1i64
        } else {
            target[i - (out_rank - target.len())]
        };
        if !(0..=i32::MAX as i64).contains(&d_t) {
            return Err(format!(
                "Expand: target dimension {d_t} must be in [0, i32::MAX]"
            ));
        }
        let d_out = if d_in == 1 {
            d_t
        } else if d_t == 1 || d_in == d_t {
            d_in
        } else {
            return Err(format!(
                "Expand: input dimension {d_in} cannot broadcast to target dimension {d_t}"
            ));
        };
        result[i] = dim_i32(d_out)?;
    }
    let r = ctx.emit(|res, s| unsafe {
        mlx::mlx_broadcast_to(res, data, result.as_ptr(), result.len(), s)
    })?;
    // Broadcast is a zero-copy stride-0 view; contiguous is deferred to the output boundary.
    ctx.bind(&n.outputs[0], r);
    Ok(())
}

fn checked_slice_axes(
    rank: usize,
    starts: &[i64],
    ends: &[i64],
    axes: &[i64],
    steps: &[i64],
) -> Result<Vec<usize>, MlxError> {
    if starts.len() != ends.len() || starts.len() != axes.len() || starts.len() != steps.len() {
        return Err("MLX Slice: starts, ends, axes and steps must have equal lengths".to_string());
    }
    if steps.contains(&0) {
        return Err("MLX Slice: step must not be zero".to_string());
    }
    let mut seen = vec![false; rank];
    let mut normalized = Vec::with_capacity(axes.len());
    for &axis in axes {
        if axis < -(rank as i64) || axis >= rank as i64 {
            return Err(format!(
                "MLX Slice: axis {axis} is out of range for rank {rank}"
            ));
        }
        let axis = norm_axis(axis, rank as i32) as usize;
        if seen[axis] {
            return Err(format!("MLX Slice: duplicate axis {axis}"));
        }
        seen[axis] = true;
        normalized.push(axis);
    }
    Ok(normalized)
}

fn slice_list(
    ctx: &mut TranslationContext,
    n: &NodeDesc,
    index: usize,
    attr: &str,
) -> Result<(bool, Vec<i64>), MlxError> {
    if present(n, index) && ctx.tensor_shape(&n.inputs[index])?.len() != 1 {
        return Err(format!("MLX Slice: {attr} must be a rank-one tensor"));
    }
    read_list_input_or_attr(ctx, n, index, attr)
}

fn slice_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let data = ctx.resolve(&n.inputs[0])?;
    let shape = ctx.shape_of(data);
    let rank = shape.len();
    if rank == 0 {
        return Err("MLX Slice: cannot slice a scalar".to_string());
    }
    let (have_starts, starts) = slice_list(ctx, n, 1, "starts")?;
    let (have_ends, ends) = slice_list(ctx, n, 2, "ends")?;
    if !have_starts || !have_ends {
        return Err("MLX Slice: starts and ends are required".to_string());
    }
    let (have_axes, mut axes) = slice_list(ctx, n, 3, "axes")?;
    if !have_axes {
        axes = (0..starts.len() as i64).collect();
    }
    let (have_steps, mut steps) = slice_list(ctx, n, 4, "steps")?;
    if !have_steps {
        steps = vec![1; starts.len()];
    }
    let axes = checked_slice_axes(rank, &starts, &ends, &axes, &steps)?;
    if starts.is_empty() {
        ctx.bind(&n.outputs[0], data);
        return Ok(());
    }

    let mut start = vec![0i32; rank];
    let mut stop = shape.clone();
    let mut stride = vec![1i32; rank];
    let mut reverse_axes: Vec<(usize, i64, i64, i64)> = Vec::new();
    for i in 0..starts.len() {
        let ax = axes[i];
        let dim = shape[ax] as i64;
        // A larger magnitude selects at most one element, so clip in integer arithmetic.
        let step = steps[i].clamp(-dim.max(1), dim.max(1));
        let (s, e) = if step > 0 {
            let s = if starts[i] < 0 {
                starts[i].saturating_add(dim)
            } else {
                starts[i]
            };
            let e = if ends[i] < 0 {
                ends[i].saturating_add(dim)
            } else {
                ends[i]
            };
            (clamp_i64(s, 0, dim), clamp_i64(e, 0, dim))
        } else {
            let s = if starts[i] < 0 {
                starts[i].saturating_add(dim)
            } else {
                starts[i]
            };
            let e = if ends[i] < 0 {
                ends[i].saturating_add(dim)
            } else {
                ends[i]
            };
            (clamp_i64(s, 0, dim - 1), clamp_i64(e, -1, dim - 1))
        };
        start[ax] = s as i32;
        stop[ax] = e as i32;
        if step > 0 {
            stride[ax] = step as i32;
        } else if s <= e {
            start[ax] = 0;
            stop[ax] = 0;
        } else {
            // MLX's slice wrapper re-normalizes a stop of -1, so it cannot represent ONNX's
            // reverse sentinel directly. Apply all forward axes as one view, then gather an
            // explicit descending index vector for each reverse axis.
            start[ax] = 0;
            stop[ax] = shape[ax];
            reverse_axes.push((ax, s, e, step));
        }
    }
    let mut r = ctx.emit(|res, s| unsafe {
        mlx::mlx_slice(
            res,
            data,
            start.as_ptr(),
            start.len(),
            stop.as_ptr(),
            stop.len(),
            stride.as_ptr(),
            stride.len(),
            s,
        )
    })?;
    for (axis, s, e, step) in reverse_axes {
        let indices = ctx.arange(s as f64, e as f64, step as f64, mlx::mlx_dtype__MLX_INT32)?;
        r = ctx.emit(|res, stream| unsafe {
            mlx::mlx_take_axis(res, r, indices, axis as i32, stream)
        })?;
    }
    // Forward Slice is a zero-copy strided view; reverse axes use explicit gathers above.
    // Contiguous materialization is otherwise deferred to the output boundary.
    ctx.bind(&n.outputs[0], r);
    Ok(())
}

fn checked_split_sizes(
    extent: Option<i64>,
    count: usize,
    sizes: Option<&[i64]>,
    num_outputs: Option<i64>,
) -> Result<Option<Vec<i64>>, MlxError> {
    if count == 0 || count > i32::MAX as usize {
        return Err("MLX Split: output count must be positive and fit int32".to_string());
    }
    if let Some(num_outputs) = num_outputs
        && (num_outputs <= 0 || num_outputs != count as i64 || sizes.is_some())
    {
        return Err(
            "MLX Split: num_outputs must match the output count and cannot accompany split sizes"
                .to_string(),
        );
    }
    if let Some(sizes) = sizes {
        if sizes.len() != count {
            return Err("MLX Split: split count must equal the output count".to_string());
        }
        let mut total = 0i64;
        for &size in sizes {
            if size < 0 {
                return Err("MLX Split: split sizes must be non-negative".to_string());
            }
            total = total
                .checked_add(size)
                .ok_or("MLX Split: split sizes sum overflows int64")?;
        }
        if extent.is_some_and(|extent| total != extent) {
            return Err("MLX Split: split sizes sum must equal the axis extent".to_string());
        }
        return Ok(Some(sizes.to_vec()));
    }
    let Some(extent) = extent else {
        return Ok(None);
    };
    let count = count as i64;
    let chunk = if num_outputs.is_some() {
        extent / count + i64::from(extent % count != 0)
    } else {
        if extent % count != 0 {
            return Err("MLX Split: equal split requires an evenly divisible axis".to_string());
        }
        extent / count
    };
    let last = extent - chunk * (count - 1);
    if last < 0 {
        return Err("MLX Split: num_outputs produces a negative final chunk".to_string());
    }
    let mut sizes = vec![chunk; count as usize];
    sizes[count as usize - 1] = last;
    Ok(Some(sizes))
}

fn split_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let data = ctx.resolve(&n.inputs[0])?;
    let shape = ctx.shape_of(data);
    let rank = shape.len() as i64;
    let axis = int_attr(n, "axis", 0);
    if axis < -rank || axis >= rank {
        return Err(format!(
            "MLX Split: axis {axis} is out of range for rank {rank}"
        ));
    }
    let axis = norm_axis(axis, rank as i32);
    if present(n, 1) && ctx.tensor_shape(&n.inputs[1])?.len() != 1 {
        return Err("MLX Split: split sizes must have rank one".to_string());
    }
    let (have_sizes, sizes) = read_list_input_or_attr(ctx, n, 1, "split")?;
    let sizes = checked_split_sizes(
        Some(i64::from(shape[axis as usize])),
        n.outputs.len(),
        have_sizes.then_some(sizes.as_slice()),
        n.ints.get("num_outputs").copied(),
    )?
    .ok_or("MLX Split: unresolved split sizes")?;
    let mut indices = Vec::with_capacity(sizes.len() - 1);
    let mut acc = 0i32;
    for &size in &sizes[..sizes.len() - 1] {
        acc += size as i32;
        indices.push(acc);
    }
    let mut parts = VectorArray::new();
    let rc = unsafe {
        mlx::mlx_split_sections(
            parts.as_mut_ptr(),
            data,
            indices.as_ptr(),
            indices.len(),
            axis,
            ctx.stream(),
        )
    };
    if rc != 0 {
        return Err("MLX Split: mlx_split_sections failed".to_string());
    }
    if parts.size() != n.outputs.len() {
        return Err("MLX Split: backend returned an unexpected output count".to_string());
    }
    for i in 0..n.outputs.len() {
        let part = ctx.keep(parts.get(i));
        // Split parts are zero-copy strided views; contiguous is deferred to the output boundary.
        ctx.bind(&n.outputs[i], part);
    }
    Ok(())
}

/// See `misc::container_output_op`: the plugin-EP Compute ABI cannot assign a sequence to a fused
/// output slot, even though OrtApi can build a standalone sequence value.
fn split_to_sequence_op(_ctx: &mut TranslationContext, _n: &NodeDesc) -> Result<(), MlxError> {
    Err(
        "MLX SplitToSequence: sequence output cannot cross the ORT 1.29 plugin EP boundary"
            .to_string(),
    )
}

fn checked_tile_geometry(
    shape: &[i64],
    repeats_shape: &[i64],
    repeats: &[i64],
    dtype: mlx::mlx_dtype,
    symbolic: bool,
) -> Result<Vec<i64>, MlxError> {
    if repeats_shape.len() != 1 {
        return Err("MLX Tile: repeats must have rank one".to_string());
    }
    if repeats_shape[0] != repeats.len() as i64 || repeats.len() != shape.len() {
        return Err("MLX Tile: repeats length must equal data rank".to_string());
    }
    let output = shape
        .iter()
        .zip(repeats)
        .map(|(&extent, &repeat)| {
            if repeat < 0 {
                return Err("MLX Tile: repeats must be nonnegative".to_string());
            }
            if extent < 0 {
                return if symbolic {
                    Ok(if repeat == 0 { 0 } else { -1 })
                } else {
                    Err("MLX Tile: negative runtime extent".to_string())
                };
            }
            let output = extent
                .checked_mul(repeat)
                .ok_or("MLX Tile: output extent overflows int64")?;
            if output > i64::from(i32::MAX) {
                return Err("MLX Tile: output extent exceeds int32".to_string());
            }
            Ok(output)
        })
        .collect::<Result<Vec<_>, MlxError>>()?;
    if output.iter().all(|&extent| extent >= 0) {
        super::host_scan::layout(&output, dtype, "Tile output")?;
    }
    Ok(output)
}

fn tile_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    if ctx.native_attention_decode() && ctx.rope_dynamic() && !n.inputs[1].constant {
        // The token-at-a-time decoder repeats are [1, 1]. Keep this as graph data instead of
        // forcing a host read during the shapeless compiled trace.
        let data = ctx.resolve(&n.inputs[0])?;
        ctx.bind(&n.outputs[0], data);
        return Ok(());
    }
    if n.inputs[1].source != Src::Initializer && !n.inputs[1].constant && ctx.is_capturing() {
        return Err("MLX Tile: live repeats cannot be captured".to_string());
    }
    let shape = ctx.tensor_shape(&n.inputs[0])?;
    let repeats_shape = ctx.tensor_shape(&n.inputs[1])?;
    checked_tile_repeats_metadata(&repeats_shape, shape.len())?;
    let repeats_int64 = if n.inputs[1].source == Src::Intermediate {
        let repeats = ctx.resolve(&n.inputs[1])?;
        ctx.dtype_of(repeats) == mlx::mlx_dtype__MLX_INT64
    } else {
        ctx.raw_host(&n.inputs[1])?.dtype
            == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64
    };
    if !repeats_int64 {
        return Err("MLX Tile: repeats must be int64".to_string());
    }
    let repeats = read_ints(ctx, n, 1)?;
    if repeats.len() as i64 != repeats_shape[0] {
        return Err("MLX Tile: repeats data length disagrees with metadata".to_string());
    }
    let dtype = mlx_dtype_from_onnx(n.outputs[0].otype);
    let output = checked_tile_geometry(&shape, &repeats_shape, &repeats, dtype, false)?;
    if output.contains(&0) {
        let output: Vec<_> = output
            .iter()
            .map(|&extent| dim_i32(extent))
            .collect::<Result<_, _>>()?;
        let result = ctx.zeros(&output, dtype)?;
        ctx.bind(&n.outputs[0], result);
        return Ok(());
    }
    super::host_scan::layout(&shape, dtype, "Tile input")?;
    let reps: Vec<i32> = repeats
        .iter()
        .map(|&repeat| dim_i32(repeat))
        .collect::<Result<_, _>>()?;
    let data = ctx.resolve(&n.inputs[0])?;
    if shape.is_empty() {
        ctx.bind(&n.outputs[0], data);
        return Ok(());
    }
    let r = ctx.emit(|res, s| unsafe { mlx::mlx_tile(res, data, reps.as_ptr(), reps.len(), s) })?;
    ctx.bind(&n.outputs[0], r);
    Ok(())
}

fn checked_tile_repeats_metadata(shape: &[i64], data_rank: usize) -> Result<(), MlxError> {
    if shape.len() != 1 {
        return Err("MLX Tile: repeats must have rank one".to_string());
    }
    if usize::try_from(shape[0]).ok() != Some(data_rank) {
        return Err("MLX Tile: repeats length must equal data rank".to_string());
    }
    Ok(())
}

struct PadGeometry {
    axes: Vec<i32>,
    output: Vec<i32>,
    start: Vec<i32>,
    stop: Vec<i32>,
    low: Vec<i32>,
    high: Vec<i32>,
    empty_source: bool,
    elements: usize,
}

fn checked_pad_axes(rank: usize, axes: &[i64], pad_count: usize) -> Result<Vec<i32>, MlxError> {
    if axes.len().checked_mul(2) != Some(pad_count) {
        return Err("MLX Pad: pads length must be twice the axes length".to_string());
    }
    let rank = i32::try_from(rank).map_err(|_| "MLX Pad: rank exceeds native int32 geometry")?;
    let mut seen = vec![false; rank as usize];
    let mut normalized = Vec::with_capacity(axes.len());
    for &axis in axes {
        if axis < -i64::from(rank) || axis >= i64::from(rank) {
            return Err(format!(
                "MLX Pad: axis {axis} is out of range for rank {rank}"
            ));
        }
        let axis = norm_axis(axis, rank);
        if seen[axis as usize] {
            return Err(format!("MLX Pad: duplicate axis {axis}"));
        }
        seen[axis as usize] = true;
        normalized.push(axis);
    }
    Ok(normalized)
}

fn checked_pad_geometry(
    shape: &[i64],
    axes: &[i64],
    pads: &[i64],
    mode: &str,
) -> Result<PadGeometry, MlxError> {
    if !matches!(mode, "constant" | "reflect" | "edge" | "wrap") {
        return Err(format!("MLX Pad: unsupported mode {mode:?}"));
    }
    let axes = checked_pad_axes(shape.len(), axes, pads.len())?;
    let axis_count = axes.len();
    let dimensions: Vec<i32> = shape
        .iter()
        .map(|&dimension| {
            if !(0..=i64::from(i32::MAX)).contains(&dimension) {
                return Err("MLX Pad: input dimensions must fit native int32 geometry".to_string());
            }
            Ok(dimension as i32)
        })
        .collect::<Result<_, _>>()?;
    let mut geometry = PadGeometry {
        axes,
        output: dimensions.clone(),
        start: vec![0; shape.len()],
        stop: dimensions,
        low: Vec::with_capacity(axis_count),
        high: Vec::with_capacity(axis_count),
        empty_source: shape.contains(&0),
        elements: 0,
    };
    for (i, &axis) in geometry.axes.iter().enumerate() {
        let axis = axis as usize;
        let extent = i128::from(shape[axis]);
        let before = i128::from(pads[i]);
        let after = i128::from(pads[i + geometry.axes.len()]);
        let output = extent + before + after;
        if !(0..=i128::from(i32::MAX)).contains(&output) {
            return Err("MLX Pad: output dimensions must fit native int32 geometry".to_string());
        }
        // Preserve the original-coordinate overlap even when removal exceeds an input edge.
        let start = (-before).clamp(0, extent);
        let stop = (extent + after).clamp(0, extent);
        let copied = (stop - start).max(0);
        if mode != "constant"
            && (extent - (-before).max(0) - (-after).max(0) < 0 || (copied == 0 && output != 0))
        {
            return Err(
                "MLX Pad: nonconstant padding requires a nonempty cropped axis".to_string(),
            );
        }
        geometry.output[axis] = output as i32;
        geometry.start[axis] = start as i32;
        geometry.stop[axis] = stop as i32;
        let low = before.clamp(0, output);
        geometry.low.push(low as i32);
        geometry.high.push((output - low - copied) as i32);
        geometry.empty_source |= copied == 0;
    }
    if !geometry.output.contains(&0) {
        geometry.elements = geometry
            .output
            .iter()
            .try_fold(1usize, |count, &dimension| {
                count
                    .checked_mul(dimension as usize)
                    .ok_or("MLX Pad: output element count overflows usize")
            })?;
    }
    Ok(geometry)
}

fn pad_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let shape = ctx.tensor_shape(&n.inputs[0])?;
    let legacy = n.since_version < 11;
    let pads = if legacy {
        let attribute = if n.since_version == 1 {
            "paddings"
        } else {
            "pads"
        };
        n.int_arrays
            .get(attribute)
            .cloned()
            .ok_or_else(|| format!("MLX Pad: missing {attribute} attribute"))?
    } else {
        if ctx.tensor_shape(&n.inputs[1])?.len() != 1 {
            return Err("MLX Pad: pads must be a rank-one int64 tensor".to_string());
        }
        read_ints(ctx, n, 1)?
    };
    let axes = if present(n, 3) {
        if ctx.tensor_shape(&n.inputs[3])?.len() != 1 {
            return Err("MLX Pad: axes must be a rank-one integer tensor".to_string());
        }
        read_ints(ctx, n, 3)?
    } else {
        (0..shape.len() as i64).collect()
    };
    if present(n, 2) && !ctx.tensor_shape(&n.inputs[2])?.is_empty() {
        return Err("MLX Pad: constant_value must be a scalar".to_string());
    }
    let mode = n
        .strings
        .get("mode")
        .map(String::as_str)
        .unwrap_or("constant");
    let geometry = checked_pad_geometry(&shape, &axes, &pads, mode)?;
    let dt = mlx_dtype_from_onnx(n.outputs[0].otype);
    if geometry.output.contains(&0) {
        let result = ctx.zeros(&geometry.output, dt)?;
        ctx.bind(&n.outputs[0], result);
        return Ok(());
    }
    let pad_value = if mode == "constant" && present(n, 2) {
        ctx.resolve(&n.inputs[2])?
    } else if mode == "constant" && legacy {
        ctx.scalar_f32(n.floats.get("value").copied().unwrap_or(0.0))
    } else {
        ctx.scalar_i64(0)
    };
    let pad_value = ctx.astype(pad_value, dt)?;
    geometry
        .elements
        .checked_mul(unsafe { mlx::mlx_array_itemsize(pad_value) })
        .ok_or("MLX Pad: output byte count overflows usize")?;
    if geometry.empty_source {
        let result = ctx.emit(|res, s| unsafe {
            mlx::mlx_broadcast_to(
                res,
                pad_value,
                geometry.output.as_ptr(),
                geometry.output.len(),
                s,
            )
        })?;
        ctx.bind(&n.outputs[0], result);
        return Ok(());
    }
    let data = ctx.resolve(&n.inputs[0])?;
    let data = if geometry.start.iter().any(|&start| start != 0)
        || geometry
            .stop
            .iter()
            .zip(&shape)
            .any(|(&stop, &extent)| i64::from(stop) != extent)
    {
        ctx.slice(data, &geometry.start, &geometry.stop)?
    } else {
        data
    };
    if geometry
        .low
        .iter()
        .chain(&geometry.high)
        .all(|&pad| pad == 0)
    {
        ctx.bind(&n.outputs[0], data);
        return Ok(());
    }
    let PadGeometry {
        axes: ax,
        low,
        high,
        ..
    } = geometry;
    let r = match mode {
        "constant" => ctx.emit(|res, s| unsafe {
            mlx::mlx_pad(
                res,
                data,
                ax.as_ptr(),
                ax.len(),
                low.as_ptr(),
                low.len(),
                high.as_ptr(),
                high.len(),
                pad_value,
                c"constant".as_ptr(),
                s,
            )
        })?,
        "reflect" => reflect_pad(ctx, data, &ax, &low, &high, pad_value)?,
        // MLX's built-in edge path aborts for multi-axis input and it has no wrap mode. Express
        // both through stable slice/broadcast/take primitives instead.
        "edge" => edge_pad(ctx, data, &ax, &low, &high)?,
        "wrap" => wrap_pad(ctx, data, &ax, &low, &high)?,
        mode => return Err(format!("Pad: unsupported mode {mode:?}")),
    };
    ctx.bind(&n.outputs[0], r);
    Ok(())
}

fn reflect_pad(
    ctx: &mut TranslationContext,
    data: mlx::mlx_array,
    axes: &[i32],
    low: &[i32],
    high: &[i32],
    pad_value: mlx::mlx_array,
) -> Result<mlx::mlx_array, MlxError> {
    let mut result = data;
    for ((&axis, &before), &after) in axes.iter().zip(low).zip(high) {
        if before == 0 && after == 0 {
            continue;
        }
        let extent = ctx.dim(result, axis);
        if extent == 1 {
            let mut shape = ctx.shape_of(result);
            shape[axis as usize] = extent + before + after;
            result = ctx.emit(|res, s| unsafe {
                mlx::mlx_broadcast_to(res, result, shape.as_ptr(), shape.len(), s)
            })?;
        } else if before < extent && after < extent {
            result = ctx.emit(|res, s| unsafe {
                mlx::mlx_pad(
                    res,
                    result,
                    &axis,
                    1,
                    &before,
                    1,
                    &after,
                    1,
                    pad_value,
                    c"reflect".as_ptr(),
                    s,
                )
            })?;
        } else {
            let period = (extent - 1)
                .checked_mul(2)
                .ok_or("MLX Pad: reflection period exceeds native int32 indexing")?;
            let indices = ctx.arange(
                -f64::from(before),
                f64::from(extent + after),
                1.0,
                mlx::mlx_dtype__MLX_INT32,
            )?;
            let period = ctx.scalar_i32(period);
            let indices = ctx.binary(mlx::mlx_remainder, indices, period)?;
            let reflected = ctx.binary(mlx::mlx_subtract, period, indices)?;
            let indices = ctx.binary(mlx::mlx_minimum, indices, reflected)?;
            result =
                ctx.emit(|res, s| unsafe { mlx::mlx_take_axis(res, result, indices, axis, s) })?;
        }
    }
    Ok(result)
}

fn edge_pad(
    ctx: &mut TranslationContext,
    data: mlx::mlx_array,
    axes: &[i32],
    low: &[i32],
    high: &[i32],
) -> Result<mlx::mlx_array, MlxError> {
    let mut result = data;
    for ((&axis, &before), &after) in axes.iter().zip(low).zip(high) {
        if before > 0 {
            let shape = ctx.shape_of(result);
            let start = vec![0; shape.len()];
            let mut stop = shape.clone();
            stop[axis as usize] = 1;
            let first = ctx.slice(result, &start, &stop)?;
            let mut target = shape;
            target[axis as usize] = before;
            let first = ctx.emit(|res, s| unsafe {
                mlx::mlx_broadcast_to(res, first, target.as_ptr(), target.len(), s)
            })?;
            result = ctx.concat2(first, result, axis)?;
        }
        if after > 0 {
            let shape = ctx.shape_of(result);
            let mut start = vec![0; shape.len()];
            start[axis as usize] = shape[axis as usize] - 1;
            let last = ctx.slice(result, &start, &shape)?;
            let mut target = shape;
            target[axis as usize] = after;
            let last = ctx.emit(|res, s| unsafe {
                mlx::mlx_broadcast_to(res, last, target.as_ptr(), target.len(), s)
            })?;
            result = ctx.concat2(result, last, axis)?;
        }
    }
    Ok(result)
}

fn wrap_pad(
    ctx: &mut TranslationContext,
    data: mlx::mlx_array,
    axes: &[i32],
    low: &[i32],
    high: &[i32],
) -> Result<mlx::mlx_array, MlxError> {
    let mut result = data;
    for ((&axis, &before), &after) in axes.iter().zip(low).zip(high) {
        let extent = ctx.dim(result, axis);
        let indices = ctx.arange(
            -(before as f64),
            (extent + after) as f64,
            1.0,
            mlx::mlx_dtype__MLX_INT32,
        )?;
        let extent = ctx.scalar_i32(extent);
        let indices = ctx.binary(mlx::mlx_remainder, indices, extent)?;
        result = ctx.emit(|res, s| unsafe { mlx::mlx_take_axis(res, result, indices, axis, s) })?;
    }
    Ok(result)
}

fn identity_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let a = ctx.resolve(&n.inputs[0])?;
    ctx.bind(&n.outputs[0], a);
    Ok(())
}

fn integer_range_count(start: i64, limit: i64, delta: i64) -> Result<usize, MlxError> {
    if delta == 0 {
        return Err("MLX Range: delta must be nonzero".to_string());
    }
    let distance = i128::from(limit) - i128::from(start);
    let delta = i128::from(delta);
    let count = if distance == 0 || distance.signum() != delta.signum() {
        0
    } else {
        (distance.abs() - 1) / delta.abs() + 1
    };
    if count > i128::from(i32::MAX) {
        return Err("MLX Range: output length exceeds native int32 geometry".to_string());
    }
    Ok(count as usize)
}

fn range_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    if ctx.native_attention_decode()
        && ctx.rope_dynamic()
        && n.inputs.iter().any(|input| !input.constant)
    {
        // Decode forms Range(start, start + 1, 1). Its one value is the dynamic start itself.
        let start = ctx.resolve(&n.inputs[0])?;
        let start = ctx.reshape(start, &[1])?;
        ctx.bind(&n.outputs[0], start);
        return Ok(());
    }
    if n.outputs[0].otype == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64 {
        let mut bounds = [0i64; 3];
        for (i, value) in bounds.iter_mut().enumerate() {
            if !ctx.tensor_shape(&n.inputs[i])?.is_empty() {
                return Err("MLX Range: inputs must be scalars".to_string());
            }
            let values = ctx.read_ints_eval(&n.inputs[i])?;
            if values.len() != 1 {
                return Err("MLX Range: expected one integer per scalar".to_string());
            }
            *value = values[0];
        }
        let [start, limit, delta] = bounds;
        let count = integer_range_count(start, limit, delta)?;
        let mut values = Vec::new();
        values
            .try_reserve_exact(count)
            .map_err(|error| format!("MLX Range: cannot allocate integer output: {error}"))?;
        for i in 0..count {
            let value = i128::from(start) + (i as i128) * i128::from(delta);
            values
                .push(i64::try_from(value).map_err(|_| "MLX Range: integer output exceeds int64")?);
        }
        let out = ctx.from_host_i64(&values, &[count as i32]);
        ctx.bind(&n.outputs[0], out);
        return Ok(());
    }
    let start = read_range_scalar(ctx, n, 0)?;
    let limit = read_range_scalar(ctx, n, 1)?;
    let delta = read_range_scalar(ctx, n, 2)?;
    let dt = mlx_dtype_from_onnx(n.outputs[0].otype);
    let double_stash = (dt == mlx::mlx_dtype__MLX_FLOAT16 || dt == mlx::mlx_dtype__MLX_BFLOAT16)
        && n.ints.get("stash_type").copied().unwrap_or(1) == 11;
    if double_stash {
        let count = ((limit - start) / delta).ceil().max(0.0) as usize;
        let values: Vec<u16> = (0..count)
            .map(|i| {
                let value = start + i as f64 * delta;
                if dt == mlx::mlx_dtype__MLX_FLOAT16 {
                    f16::from_f64(value).to_bits()
                } else {
                    bf16::from_f64(value).to_bits()
                }
            })
            .collect();
        let array = Array::from_data(
            values.as_ptr() as *const std::os::raw::c_void,
            &[count as i32],
            dt,
        );
        array.eval();
        let out = ctx.keep(array);
        ctx.bind(&n.outputs[0], out);
        return Ok(());
    }
    let compute_dt = if dt == mlx::mlx_dtype__MLX_FLOAT16 || dt == mlx::mlx_dtype__MLX_BFLOAT16 {
        mlx::mlx_dtype__MLX_FLOAT32
    } else {
        dt
    };
    let r =
        ctx.emit(|res, s| unsafe { mlx::mlx_arange(res, start, limit, delta, compute_dt, s) })?;
    let r = if compute_dt == dt {
        r
    } else {
        ctx.astype(r, dt)?
    };
    ctx.bind(&n.outputs[0], r);
    Ok(())
}

fn read_range_scalar(
    ctx: &mut TranslationContext,
    n: &NodeDesc,
    i: usize,
) -> Result<f64, MlxError> {
    if n.inputs[i].source == Src::Intermediate {
        return ctx
            .read_ints_eval(&n.inputs[i])?
            .first()
            .map(|&value| value as f64)
            .ok_or_else(|| "Range expected a scalar input".to_string());
    }
    let h = ctx.raw_host(&n.inputs[i])?;
    if h.count != 1 || h.data.is_null() {
        return Err("Range expected a scalar initializer".to_string());
    }
    match h.dtype {
        t if t == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16 => {
            Ok(unsafe { *(h.data as *const i16) } as f64)
        }
        t if t == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32 => {
            Ok(unsafe { *(h.data as *const i32) } as f64)
        }
        t if t == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64 => {
            Ok(unsafe { *(h.data as *const i64) } as f64)
        }
        t if t == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT => {
            Ok(unsafe { *(h.data as *const f32) } as f64)
        }
        t if t == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16 => {
            Ok(f16::from_bits(unsafe { *(h.data as *const u16) }).to_f64())
        }
        t if t == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16 => {
            Ok(bf16::from_bits(unsafe { *(h.data as *const u16) }).to_f64())
        }
        _ => Err("Range initializer dtype is not supported".to_string()),
    }
}

fn shape_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let input_shape = ctx.tensor_shape(&n.inputs[0])?;
    let rank = input_shape.len() as i64;
    let start_attr = int_attr(n, "start", 0);
    let end_attr = int_attr(n, "end", rank);
    let (start, end) = shape_interval(rank, start_attr, end_attr);
    let result: Vec<i64> = (start..end).map(|i| input_shape[i as usize]).collect();
    let length = i32::try_from(result.len())
        .map_err(|_| "MLX Shape: output length exceeds native i32 shape range".to_string())?;
    let out = ctx.from_host_i64(&result, &[length]);
    ctx.bind(&n.outputs[0], out);
    Ok(())
}

fn size_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let shape = ctx.tensor_shape(&n.inputs[0])?;
    if shape.iter().any(|&dim| dim < 0) {
        return Err("MLX Size: runtime tensor has a negative dimension".to_string());
    }
    let size = if shape.contains(&0) {
        0
    } else {
        shape.iter().try_fold(1i64, |size, &dim| {
            size.checked_mul(dim)
                .ok_or_else(|| "MLX Size: element count exceeds int64".to_string())
        })?
    };
    let out = ctx.from_host_i64(&[size], &[]);
    ctx.bind(&n.outputs[0], out);
    Ok(())
}

fn shape_interval(rank: i64, start: i64, end: i64) -> (i64, i64) {
    let mut s = if start < 0 { start + rank } else { start };
    let mut e = if end < 0 { end + rank } else { end };
    s = clamp_i64(s, 0, rank);
    e = clamp_i64(e, 0, rank);
    if e < s {
        e = s;
    }
    (s, e)
}

fn constant_of_shape_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let shape_i64 = read_ints(ctx, n, 0)?;
    let s: Vec<i32> = shape_i64
        .iter()
        .map(|&dimension| {
            if !(0..=i32::MAX as i64).contains(&dimension) {
                return Err(format!(
                    "ConstantOfShape: dimension {dimension} must be in [0, i32::MAX]"
                ));
            }
            dim_i32(dimension)
        })
        .collect::<Result<_, _>>()?;
    let out_type = n.outputs[0].otype;
    let mlx_type = crate::engine::mlx_dtype_from_onnx(out_type);

    // ONNX: `value` is a one-element tensor supplying the fill value and the output dtype; absent,
    // the output is fp32 zeros. Read the attribute when present so the explicit-value form is not
    // stranded on CPU (each such node is a partition boundary).
    let fill = n
        .tensors
        .get("value")
        .map(|t| {
            t.single_element_bytes()
                .map(|b| (b.to_vec(), t.dtype))
                .ok_or_else(|| {
                    "ConstantOfShape: value must contain exactly one element".to_string()
                })
        })
        .transpose()?;
    let r = match fill {
        Some((bytes, value_dtype)) => {
            // ONNX takes the output dtype from `value`, so build the scalar in that dtype and let
            // mlx_full carry it, with no widening round-trip. Note the scalar must be built in its
            // own dtype rather than staged through float64 as a common transport: an fp64 array is
            // only evaluable on an MLX CPU stream, and a plan gets that stream only when
            // GetCapability coloured it fp64, so routing an fp32 fill through float64 would abort
            // the GPU-stream plan it lands in.
            let scalar =
                ctx.scalar_from_bytes(&bytes, crate::engine::mlx_dtype_from_onnx(value_dtype));
            ctx.emit(|res, st| unsafe {
                mlx::mlx_full(res, s.as_ptr(), s.len(), scalar, mlx_type, st)
            })?
        }
        None => ctx.zeros(&s, mlx_type)?,
    };
    ctx.bind(&n.outputs[0], r);
    Ok(())
}

fn where_handler(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let cond = ctx.resolve(&n.inputs[0])?;
    let x = ctx.resolve(&n.inputs[1])?;
    let y = ctx.resolve(&n.inputs[2])?;
    let r = where_op(ctx, cond, x, y)?;
    ctx.bind(&n.outputs[0], r);
    Ok(())
}

fn resize_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let mut data = ctx.resolve(&n.inputs[0])?;
    let input_shape = ctx.shape_of(data);
    let rank = input_shape.len();
    let legacy_upsample = n.op_type == "Upsample";
    let output_dtype = mlx_dtype_from_onnx(n.outputs[0].otype);
    let mode = n
        .strings
        .get("mode")
        .map(String::as_str)
        .unwrap_or("nearest");
    let coordinate_mode = if legacy_upsample {
        "asymmetric"
    } else {
        n.strings
            .get("coordinate_transformation_mode")
            .map(String::as_str)
            .unwrap_or("half_pixel")
    };
    let nearest_mode = if legacy_upsample {
        "floor"
    } else {
        n.strings
            .get("nearest_mode")
            .map(String::as_str)
            .unwrap_or("round_prefer_floor")
    };

    let (output_lengths, scales): (Vec<i32>, Vec<f64>) = if !legacy_upsample && present(n, 3) {
        let sizes = read_ints(ctx, n, 3)?;
        (
            sizes.iter().map(|&size| size as i32).collect(),
            sizes
                .iter()
                .enumerate()
                .map(|(i, &size)| size as f64 / input_shape[i] as f64)
                .collect(),
        )
    } else {
        let scales_input = if legacy_upsample { 1 } else { 2 };
        let scales_host = ctx.raw_host(&n.inputs[scales_input])?;
        if scales_host.data.is_null() || scales_host.count != rank {
            return Err(format!(
                "MLX {} requires constant float32 scales",
                n.op_type
            ));
        }
        let scales = unsafe {
            std::slice::from_raw_parts(scales_host.data as *const f32, scales_host.count)
        };
        (
            scales
                .iter()
                .enumerate()
                .map(|(i, &scale)| (scale as f64 * input_shape[i] as f64).floor() as i32)
                .collect(),
            scales.iter().map(|&scale| scale as f64).collect(),
        )
    };

    let restore_dtype = ctx.dtype_of(data);
    let use_f32 = mode == "linear" || restore_dtype == mlx::mlx_dtype__MLX_BFLOAT16;
    if use_f32 {
        data = ctx.astype(data, mlx::mlx_dtype__MLX_FLOAT32)?;
    }
    for axis in 0..rank {
        let input_len = input_shape[axis] as i64;
        let output_len = output_lengths[axis] as i64;
        // Equal rounded extents alone do not make fractional-scale coordinates an identity.
        if input_len == output_len
            && (scales[axis] == 1.0 || coordinate_mode == "align_corners" || input_len <= 1)
        {
            continue;
        }
        if mode == "nearest" {
            let indices: Vec<i32> = (0..output_len)
                .map(|i| {
                    resize_nearest_index(
                        nearest_mode,
                        resize_src_coord(coordinate_mode, i, scales[axis], input_len, output_len),
                        input_len,
                    )
                })
                .collect();
            let index = ctx.keep(Array::from_data(
                indices.as_ptr() as *const std::os::raw::c_void,
                &[output_len as i32],
                mlx::mlx_dtype__MLX_INT32,
            ));
            data =
                ctx.emit(|res, s| unsafe { mlx::mlx_take_axis(res, data, index, axis as i32, s) })?;
            continue;
        }

        let mut lower = Vec::with_capacity(output_len as usize);
        let mut upper = Vec::with_capacity(output_len as usize);
        let mut lower_weight = Vec::with_capacity(output_len as usize);
        let mut upper_weight = Vec::with_capacity(output_len as usize);
        for i in 0..output_len {
            let source = resize_src_coord(coordinate_mode, i, scales[axis], input_len, output_len);
            let floor = source.floor();
            lower.push(clamp_i64(floor as i64, 0, input_len - 1) as i32);
            upper.push(clamp_i64(floor as i64 + 1, 0, input_len - 1) as i32);
            upper_weight.push((source - floor) as f32);
            lower_weight.push((1.0 - (source - floor)) as f32);
        }
        let lower_index = ctx.keep(Array::from_data(
            lower.as_ptr() as *const std::os::raw::c_void,
            &[output_len as i32],
            mlx::mlx_dtype__MLX_INT32,
        ));
        let upper_index = ctx.keep(Array::from_data(
            upper.as_ptr() as *const std::os::raw::c_void,
            &[output_len as i32],
            mlx::mlx_dtype__MLX_INT32,
        ));
        let mut weight_shape = vec![1i32; rank];
        weight_shape[axis] = output_len as i32;
        let lw = ctx.keep(Array::from_data(
            lower_weight.as_ptr() as *const std::os::raw::c_void,
            &weight_shape,
            mlx::mlx_dtype__MLX_FLOAT32,
        ));
        let uw = ctx.keep(Array::from_data(
            upper_weight.as_ptr() as *const std::os::raw::c_void,
            &weight_shape,
            mlx::mlx_dtype__MLX_FLOAT32,
        ));
        let lo = ctx
            .emit(|res, s| unsafe { mlx::mlx_take_axis(res, data, lower_index, axis as i32, s) })?;
        let hi = ctx
            .emit(|res, s| unsafe { mlx::mlx_take_axis(res, data, upper_index, axis as i32, s) })?;
        let lo = ctx.mul(lo, lw)?;
        let hi = ctx.mul(hi, uw)?;
        data = ctx.add(lo, hi)?;
    }
    if use_f32 {
        data = ctx.astype(data, output_dtype)?;
    }
    let data = ctx.contiguous(data)?;
    ctx.bind(&n.outputs[0], data);
    Ok(())
}

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

fn movable_io(
    node: &NodeView,
) -> Option<(
    ort::ONNXTensorElementDataType,
    ort::ONNXTensorElementDataType,
)> {
    match (node.input_info(0), node.output_info(0)) {
        (Some(i), Some(o)) => Some((i.dtype, o.dtype)),
        _ => None,
    }
}

fn gather_like_claim(node: &NodeView) -> ClaimResult {
    require!(
        node.num_inputs() == 2 && node.num_outputs() == 1,
        "Gather/GatherElements expects 2 inputs and 1 output, got {}in/{}out",
        node.num_inputs(),
        node.num_outputs()
    );
    let (data, out) = match movable_io(node) {
        Some(v) => v,
        None => deny!("missing tensor type/shape info on the data input or output"),
    };
    let idx = match node.input_info(1) {
        Some(i) => i.dtype,
        None => deny!("missing tensor type/shape info on the indices input"),
    };
    require!(
        is_movable(data),
        "data dtype {} is not a movable dtype supported on GPU",
        crate::registry::ort_dtype_name(data)
    );
    require!(
        out == data,
        "output dtype {} must match data dtype {}",
        crate::registry::ort_dtype_name(out),
        crate::registry::ort_dtype_name(data)
    );
    require!(
        is_int_index(idx),
        "indices dtype {} must be int32/int64",
        crate::registry::ort_dtype_name(idx)
    );
    Ok(())
}

fn scatter_elements_claim(node: &NodeView) -> ClaimResult {
    require!(
        node.num_inputs() == 3 && node.num_outputs() == 1,
        "ScatterElements expects 3 inputs and 1 output, got {}in/{}out",
        node.num_inputs(),
        node.num_outputs()
    );
    require!(
        node.string_attr("reduction", "none") == "none",
        "ScatterElements: only reduction=none is claimed (add/mul/min/max reductions stay on CPU)"
    );
    let (data, indices, updates, out) = match (
        node.input_info(0),
        node.input_info(1),
        node.input_info(2),
        node.output_info(0),
    ) {
        (Some(a), Some(b), Some(c), Some(d)) => (a, b, c, d),
        _ => deny!("missing tensor type/shape info on an input or the output"),
    };
    // mlx_put_along_axis's GPU kernel aborts on int64 payloads → keep to MLX float dtypes.
    require!(
        is_mlx_float(data.dtype),
        "data dtype {} unsupported: mlx_put_along_axis's GPU kernel needs an MLX float (fp32/fp16/bf16)",
        crate::registry::ort_dtype_name(data.dtype)
    );
    require!(
        is_int_index(indices.dtype),
        "indices dtype {} must be int32/int64",
        crate::registry::ort_dtype_name(indices.dtype)
    );
    require!(
        updates.dtype == data.dtype && out.dtype == data.dtype,
        "updates/output dtype must match data dtype {} (got updates {}, out {})",
        crate::registry::ort_dtype_name(data.dtype),
        crate::registry::ort_dtype_name(updates.dtype),
        crate::registry::ort_dtype_name(out.dtype)
    );
    require!(
        !data.shape.is_empty(),
        "data must have rank >= 1 (scalar data is unsupported)"
    );
    require!(
        compatible_shape(&indices.shape, &updates.shape),
        "indices shape {:?} must equal updates shape {:?}",
        indices.shape,
        updates.shape
    );
    require!(
        indices.shape.len() == data.shape.len(),
        "indices rank {} must equal data rank {}",
        indices.shape.len(),
        data.shape.len()
    );
    require!(
        compatible_shape(&out.shape, &data.shape),
        "output shape {:?} must equal data shape {:?}",
        out.shape,
        data.shape
    );
    let rank = data.shape.len() as i64;
    let axis = node.int_attr("axis", 0);
    require!(
        axis >= -rank && axis < rank,
        "axis {axis} is out of range for rank {rank}"
    );
    let ax = norm_axis(axis, rank as i32) as usize;
    for i in 0..data.shape.len() {
        require!(
            i == ax
                || indices.shape[i] < 0
                || data.shape[i] < 0
                || indices.shape[i] <= data.shape[i],
            "on non-axis dim {i}, indices extent {} must be <= data extent {}",
            indices.shape[i],
            data.shape[i]
        );
    }
    Ok(())
}

fn center_crop_pad_claim(node: &NodeView) -> ClaimResult {
    require!(
        node.num_inputs() == 2 && node.num_outputs() == 1,
        "CenterCropPad expects 2 inputs and 1 output, got {}in/{}out",
        node.num_inputs(),
        node.num_outputs()
    );
    let (data, output) = match (node.input_info(0), node.output_info(0)) {
        (Some(data), Some(output)) => (data, output),
        _ => deny!("CenterCropPad: missing tensor type/shape info"),
    };
    require!(
        is_movable(data.dtype) && output.dtype == data.dtype,
        "CenterCropPad: input/output must share a movable dtype"
    );
    require!(
        !data.shape.is_empty()
            && data.shape.iter().all(|&d| d >= 0)
            && output.shape.iter().all(|&d| d >= 0),
        "CenterCropPad: input/output shapes must be fully static"
    );
    let target = match node.read_const_int64(1) {
        Some(target) => target,
        None => deny!("CenterCropPad: `shape` must be a constant int64 initializer"),
    };
    let rank = data.shape.len() as i64;
    let (have_axes, axes) = node.ints_attr("axes");
    let axes = if have_axes { axes } else { (0..rank).collect() };
    require!(
        target.len() == axes.len(),
        "CenterCropPad: shape length {} must equal axes length {}",
        target.len(),
        axes.len()
    );
    require!(
        output.shape.len() == data.shape.len(),
        "CenterCropPad: output rank must equal input rank"
    );
    let mut expected = data.shape.clone();
    let mut seen = vec![false; rank as usize];
    for (&axis, &size) in axes.iter().zip(&target) {
        require!(
            axis >= -rank && axis < rank,
            "CenterCropPad: axis {axis} is out of range for rank {rank}"
        );
        require!(
            size >= 0 && size <= i32::MAX as i64,
            "CenterCropPad: target sizes must be in [0, i32::MAX]"
        );
        let axis = norm_axis(axis, rank as i32) as usize;
        require!(!seen[axis], "CenterCropPad: axes must be unique");
        seen[axis] = true;
        expected[axis] = size;
    }
    require!(
        output.shape == expected,
        "CenterCropPad: output shape {:?} must equal expected {:?}",
        output.shape,
        expected
    );
    Ok(())
}

fn compress_claim(node: &NodeView) -> ClaimResult {
    require!(
        node.num_inputs() == 2 && node.num_outputs() == 1,
        "Compress expects 2 inputs and 1 output, got {}in/{}out",
        node.num_inputs(),
        node.num_outputs()
    );
    let (data, condition, output) =
        match (node.input_info(0), node.input_info(1), node.output_info(0)) {
            (Some(data), Some(condition), Some(output)) => (data, condition, output),
            _ => deny!("Compress: missing tensor type/shape info"),
        };
    require!(
        is_movable(data.dtype) && output.dtype == data.dtype,
        "Compress: input/output must share a movable dtype"
    );
    require!(
        condition.dtype == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL
            && condition.shape.len() == 1,
        "Compress: condition must be a rank-1 bool tensor"
    );
    require!(
        node.is_constant_initializer(1),
        "Compress: condition must be a constant initializer because the output extent is value-dependent"
    );
    require!(
        data.shape.iter().all(|&d| d >= 0) && output.shape.iter().all(|&d| d >= 0),
        "Compress: input/output shapes must be fully static"
    );
    if node.has_attr("axis") {
        let rank = data.shape.len() as i64;
        let axis = node.int_attr("axis", 0);
        require!(
            rank > 0 && axis >= -rank && axis < rank,
            "Compress: axis {axis} is out of range for rank {rank}"
        );
        require!(
            output.shape.len() == data.shape.len(),
            "Compress: output rank must equal input rank when axis is present"
        );
        let ax = norm_axis(axis, rank as i32) as usize;
        require!(
            output
                .shape
                .iter()
                .enumerate()
                .all(|(i, &d)| i == ax || d == data.shape[i]),
            "Compress: non-axis output dimensions must equal the input dimensions"
        );
    } else {
        require!(
            output.shape.len() == 1,
            "Compress without axis must produce a rank-1 output"
        );
    }
    Ok(())
}

fn depth_to_space_claim(node: &NodeView) -> ClaimResult {
    require!(
        node.num_inputs() == 1 && node.num_outputs() == 1,
        "DepthToSpace expects 1 input and 1 output"
    );
    let (data, output) = match (node.input_info(0), node.output_info(0)) {
        (Some(data), Some(output)) => (data, output),
        _ => deny!("DepthToSpace: missing tensor type/shape info"),
    };
    require!(
        is_movable(data.dtype) && output.dtype == data.dtype,
        "DepthToSpace: input/output must share a movable dtype"
    );
    require!(
        data.shape.len() == 4 && data.shape.iter().all(|&d| d >= 0),
        "DepthToSpace: only fully-static rank-4 NCHW input is supported"
    );
    let block = node.int_attr("blocksize", 0);
    require!(
        block > 0 && block <= i32::MAX as i64,
        "DepthToSpace: blocksize must be positive"
    );
    require!(
        data.shape[1] % (block * block) == 0,
        "DepthToSpace: channel extent must be divisible by blocksize squared"
    );
    let mode = node.string_attr("mode", "DCR");
    require!(
        mode == "DCR" || mode == "CRD",
        "DepthToSpace: mode must be DCR or CRD"
    );
    let expected = vec![
        data.shape[0],
        data.shape[1] / (block * block),
        data.shape[2] * block,
        data.shape[3] * block,
    ];
    require!(
        output.shape == expected,
        "DepthToSpace: output shape {:?} must equal expected {:?}",
        output.shape,
        expected
    );
    Ok(())
}

fn eye_like_claim(node: &NodeView) -> ClaimResult {
    require!(
        node.num_inputs() == 1 && node.num_outputs() == 1,
        "EyeLike expects 1 input and 1 output"
    );
    let (input, output) = match (node.input_info(0), node.output_info(0)) {
        (Some(input), Some(output)) => (input, output),
        _ => deny!("EyeLike: missing tensor type/shape info"),
    };
    require!(
        input.shape.len() == 2 && input.shape.iter().all(|&d| d >= 0),
        "EyeLike: input must be a fully-static rank-2 tensor"
    );
    require!(
        is_mlx_numeric(input.dtype) && is_mlx_numeric(output.dtype),
        "EyeLike: input/output dtypes must be MLX-compatible numeric types"
    );
    require!(
        output.shape == input.shape,
        "EyeLike: output shape must equal input shape"
    );
    let expected_dtype = node.int_attr("dtype", input.dtype as i64);
    require!(
        expected_dtype == output.dtype as i64,
        "EyeLike: output dtype must match the dtype attribute or input dtype"
    );
    let k = node.int_attr("k", 0);
    require!(
        k >= i32::MIN as i64 && k <= i32::MAX as i64,
        "EyeLike: k must fit int32"
    );
    Ok(())
}

fn reverse_sequence_claim(node: &NodeView) -> ClaimResult {
    require!(
        node.num_inputs() == 2 && node.num_outputs() == 1,
        "ReverseSequence expects 2 inputs and 1 output"
    );
    let (data, lengths, output) =
        match (node.input_info(0), node.input_info(1), node.output_info(0)) {
            (Some(data), Some(lengths), Some(output)) => (data, lengths, output),
            _ => deny!("ReverseSequence: missing tensor type/shape info"),
        };
    require!(
        is_movable(data.dtype) && output.dtype == data.dtype && output.shape == data.shape,
        "ReverseSequence: input/output must have the same movable dtype and shape"
    );
    require!(
        data.shape.len() >= 2 && data.shape.iter().all(|&d| d >= 0),
        "ReverseSequence: input must have rank >= 2 and a fully-static shape"
    );
    require!(
        lengths.dtype == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64
            && lengths.shape.len() == 1,
        "ReverseSequence: sequence_lens must be rank-1 int64"
    );
    let rank = data.shape.len() as i64;
    let batch_axis = node.int_attr("batch_axis", 1);
    let time_axis = node.int_attr("time_axis", 0);
    require!(
        batch_axis >= -rank && batch_axis < rank && time_axis >= -rank && time_axis < rank,
        "ReverseSequence: batch_axis/time_axis must be in range"
    );
    let batch_axis = norm_axis(batch_axis, rank as i32) as usize;
    let time_axis = norm_axis(time_axis, rank as i32) as usize;
    require!(
        batch_axis != time_axis,
        "ReverseSequence: batch_axis and time_axis must differ"
    );
    require!(
        lengths.shape[0] == data.shape[batch_axis],
        "ReverseSequence: sequence_lens length must equal the batch extent"
    );
    Ok(())
}

fn scatter_nd_compile_policy(reduction: &str) -> crate::registry::CompileShapeSafety {
    if reduction == "add" {
        crate::registry::CompileShapeSafety::EagerOnly {
            reason: "validates live ScatterND add indices before unchecked MLX scatter writes",
        }
    } else {
        crate::registry::CompileShapeSafety::ShapeKeyedOnly {
            reason: crate::registry::MLX_SCATTER_SHAPE_REASON,
        }
    }
}

fn scatter_nd_claim(node: &NodeView) -> ClaimResult {
    require!(
        node.num_inputs() == 3 && node.num_outputs() == 1,
        "ScatterND expects 3 inputs and 1 output"
    );
    let reduction = node.string_attr("reduction", "none");
    require!(
        reduction == "none" || reduction == "add",
        "ScatterND: supports reduction none/add"
    );
    require!(
        reduction != "add" || node.since_version() >= 16,
        "ScatterND: reduction add requires opset 16 or newer"
    );
    let (data, indices, updates, output) = match (
        node.input_info(0),
        node.input_info(1),
        node.input_info(2),
        node.output_info(0),
    ) {
        (Some(data), Some(indices), Some(updates), Some(output)) => {
            (data, indices, updates, output)
        }
        _ => deny!("ScatterND: missing tensor type/shape info"),
    };
    require!(
        is_mlx_float(data.dtype),
        "ScatterND: data must be fp32/fp16/bf16 because MLX scatter payload kernels require float"
    );
    let compatible = |left: &[i64], right: &[i64]| {
        left.len() == right.len()
            && left
                .iter()
                .zip(right)
                .all(|(&a, &b)| a >= -1 && b >= -1 && (a < 0 || b < 0 || a == b))
    };
    require!(
        updates.dtype == data.dtype
            && output.dtype == data.dtype
            && compatible(&data.shape, &output.shape),
        "ScatterND: data/output dtype or shape mismatch"
    );
    require!(
        is_int_index(indices.dtype) && !indices.shape.is_empty() && !data.shape.is_empty(),
        "ScatterND: requires integer index tuples and positive ranks"
    );
    let width = *indices.shape.last().unwrap();
    require!(
        width > 0 && width as usize <= data.shape.len(),
        "ScatterND: index tuple width must be static and within data rank"
    );
    let mut expected = indices.shape[..indices.shape.len() - 1].to_vec();
    expected.extend_from_slice(&data.shape[width as usize..]);
    require!(
        compatible(&expected, &updates.shape),
        "ScatterND: incompatible updates shape"
    );
    Ok(())
}

fn space_to_depth_claim(node: &NodeView) -> ClaimResult {
    require!(
        node.num_inputs() == 1 && node.num_outputs() == 1,
        "SpaceToDepth expects 1 input and 1 output"
    );
    let (data, output) = match (node.input_info(0), node.output_info(0)) {
        (Some(data), Some(output)) => (data, output),
        _ => deny!("SpaceToDepth: missing tensor type/shape info"),
    };
    require!(
        is_movable(data.dtype) && output.dtype == data.dtype,
        "SpaceToDepth: input/output must share a movable dtype"
    );
    require!(
        data.shape.len() == 4 && data.shape.iter().all(|&d| d >= 0),
        "SpaceToDepth: only fully-static rank-4 NCHW input is supported"
    );
    let block = node.int_attr("blocksize", 0);
    require!(
        block > 0 && block <= i32::MAX as i64,
        "SpaceToDepth: blocksize must be positive"
    );
    require!(
        data.shape[2] % block == 0 && data.shape[3] % block == 0,
        "SpaceToDepth: spatial extents must be divisible by blocksize"
    );
    let mode = node.string_attr("mode", "DCR");
    require!(
        mode == "DCR" || mode == "CRD",
        "SpaceToDepth: mode must be DCR or CRD"
    );
    let expected = vec![
        data.shape[0],
        data.shape[1] * block * block,
        data.shape[2] / block,
        data.shape[3] / block,
    ];
    require!(
        output.shape == expected,
        "SpaceToDepth: output shape {:?} must equal expected {:?}",
        output.shape,
        expected
    );
    Ok(())
}

fn concat_claim(node: &NodeView) -> ClaimResult {
    require!(
        node.num_inputs() != 0 && node.num_outputs() == 1,
        "Concat expects 1+ inputs and 1 output, got {}in/{}out",
        node.num_inputs(),
        node.num_outputs()
    );
    let out = match node.output_info(0) {
        Some(o) if is_movable(o.dtype) => o.dtype,
        Some(o) => deny!(
            "output dtype {} is not a movable dtype supported on GPU",
            crate::registry::ort_dtype_name(o.dtype)
        ),
        None => deny!("missing output tensor type/shape info"),
    };
    for i in 0..node.num_inputs() {
        match node.input_info(i) {
            Some(info) if info.dtype == out => {}
            Some(info) => deny!(
                "input[{i}] dtype {} must match output dtype {}",
                crate::registry::ort_dtype_name(info.dtype),
                crate::registry::ort_dtype_name(out)
            ),
            None => deny!("input[{i}] has no tensor type/shape info"),
        }
    }
    Ok(())
}

fn reshape_claim(node: &NodeView) -> ClaimResult {
    require!(
        node.num_inputs() == 2 && node.num_outputs() == 1,
        "Reshape expects 2 inputs and 1 output, got {}in/{}out",
        node.num_inputs(),
        node.num_outputs()
    );
    let (data, out) = match movable_io(node) {
        Some(v) => v,
        None => deny!("missing tensor type/shape info on the data input or output"),
    };
    require!(
        is_movable(data),
        "data dtype {} is not a movable dtype supported on GPU",
        crate::registry::ort_dtype_name(data)
    );
    require!(
        out == data,
        "output dtype {} must match data dtype {}",
        crate::registry::ort_dtype_name(out),
        crate::registry::ort_dtype_name(data)
    );
    require!(
        node.input_info(1).map(|i| i.dtype)
            == Some(ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64),
        "Reshape `shape` input must be int64 (runtime shape-derived values resolve at trace time; data-dependent geometry uses eager translation and cyclic fusions are split)"
    );
    Ok(())
}

fn movable_view_claim(node: &NodeView) -> ClaimResult {
    require!(
        node.num_inputs() == 1 && node.num_outputs() == 1,
        "expects 1 input and 1 output, got {}in/{}out",
        node.num_inputs(),
        node.num_outputs()
    );
    let (data, out) = match movable_io(node) {
        Some(v) => v,
        None => deny!("missing tensor type/shape info on the input or output"),
    };
    require!(
        is_movable(data),
        "input dtype {} is not a movable numeric/Boolean dtype",
        crate::registry::ort_dtype_name(data)
    );
    require!(
        out == data,
        "output dtype {} must match input dtype {}",
        crate::registry::ort_dtype_name(out),
        crate::registry::ort_dtype_name(data)
    );
    Ok(())
}

fn transpose_claim(node: &NodeView) -> ClaimResult {
    movable_view_claim(node)?;
    let info = match node.input_info(0) {
        Some(info) => info,
        None => deny!("Transpose requires tensor data"),
    };
    require!(
        info.shape.iter().all(|&dim| dim <= i64::from(i32::MAX)),
        "Transpose input dimensions must fit int32"
    );
    let (have_perm, perm) = node.ints_attr("perm");
    if node.input_rank_known(0) == Some(true) || have_perm {
        let rank = if node.input_rank_known(0) == Some(true) {
            info.shape.len()
        } else {
            perm.len()
        };
        let checked = checked_transpose_perm(rank, have_perm.then_some(perm.as_slice()));
        require!(checked.is_ok(), "{}", checked.unwrap_err());
    }
    Ok(())
}

fn transpose_compile_shape_safety(node: &NodeView) -> crate::registry::CompileShapeSafety {
    if node.input_rank_known(0) == Some(true) {
        crate::registry::CompileShapeSafety::Shapeless
    } else {
        crate::registry::CompileShapeSafety::ShapeKeyedOnly {
            reason: "a runtime input rank determines the transpose permutation",
        }
    }
}

fn rank_change_claim(node: &NodeView) -> ClaimResult {
    require!(
        (1..=2).contains(&node.num_inputs()) && node.num_outputs() == 1,
        "Unsqueeze/Squeeze expects 1 or 2 inputs and 1 output, got {}in/{}out",
        node.num_inputs(),
        node.num_outputs()
    );
    let (data, out) = match movable_io(node) {
        Some(v) => v,
        None => deny!("missing tensor type/shape info on the data input or output"),
    };
    require!(
        is_movable(data),
        "data dtype {} is not a movable numeric/Boolean dtype",
        crate::registry::ort_dtype_name(data)
    );
    require!(
        out == data,
        "output dtype {} must match data dtype {}",
        crate::registry::ort_dtype_name(out),
        crate::registry::ort_dtype_name(data)
    );
    let inserting = node.op_type() == "Unsqueeze";
    let (have_axes, axes) = if node.since_version() >= 13 {
        let have_axes = node
            .input_names()
            .get(1)
            .is_some_and(|name| !name.is_empty());
        if have_axes {
            let info = match node.input_info(1) {
                Some(info) => info,
                None => deny!("Squeeze/Unsqueeze axes must be an int64 tensor"),
            };
            require!(
                info.dtype == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64
                    && (node.input_rank_known(1) != Some(true) || info.shape.len() == 1),
                "Squeeze/Unsqueeze axes must be a one-dimensional int64 tensor"
            );
        }
        (have_axes, node.read_const_int64(1))
    } else {
        require!(
            node.num_inputs() == 1,
            "legacy Squeeze/Unsqueeze axes must be an attribute"
        );
        let (have, axes) = node.ints_attr("axes");
        (have, have.then_some(axes))
    };
    require!(!inserting || have_axes, "Unsqueeze requires axes");
    let info = match node.input_info(0) {
        Some(info) => info,
        None => deny!("Squeeze/Unsqueeze requires tensor data"),
    };
    require!(
        info.shape.iter().all(|&dim| dim <= i64::from(i32::MAX)),
        "Squeeze/Unsqueeze dimensions must fit int32"
    );
    if !inserting && node.since_version() < 13 && axes.as_deref() == Some(&[]) {
        require!(
            node.input_rank_known(0) == Some(true)
                && info.shape.iter().all(|&dim| dim >= 0 && dim != 1),
            "legacy Squeeze empty axes: shape inference and CPU/reference disagree for singleton dimensions"
        );
    }
    if node.input_rank_known(0) == Some(true) && (!have_axes || axes.is_some()) {
        let checked = checked_rank_change_axes(
            &info.shape,
            axes.as_deref(),
            inserting,
            node.since_version() >= 11,
        );
        require!(checked.is_ok(), "{}", checked.unwrap_err());
    }
    Ok(())
}

fn rank_change_compile_shape_safety(node: &NodeView) -> crate::registry::CompileShapeSafety {
    use crate::registry::CompileShapeSafety;
    if node.input_rank_known(0) == Some(true)
        && let Some(info) = node.input_info(0)
    {
        let (have_axes, axes) = if node.since_version() >= 13 {
            let have = node
                .input_names()
                .get(1)
                .is_some_and(|name| !name.is_empty());
            (have, node.read_const_int64(1))
        } else {
            let (have, axes) = node.ints_attr("axes");
            (have, have.then_some(axes))
        };
        if node.op_type() == "Unsqueeze" && axes.is_some() {
            return CompileShapeSafety::Shapeless;
        }
        if (!have_axes && info.shape.iter().all(|&dim| dim >= 0))
            || axes.as_deref().is_some_and(|axes| {
                checked_rank_change_axes(&info.shape, Some(axes), false, true)
                    .is_ok_and(|axes| axes.iter().all(|&axis| info.shape[axis as usize] == 1))
            })
        {
            return CompileShapeSafety::Shapeless;
        }
    }
    CompileShapeSafety::ShapeKeyedOnly {
        reason: "live axes and changing rank/singleton geometry require shape/value-keyed traces",
    }
}

fn flatten_claim(node: &NodeView) -> ClaimResult {
    movable_view_claim(node)?;
    let info = match node.input_info(0) {
        Some(info) => info,
        None => deny!("Flatten requires tensor data"),
    };
    require!(
        info.shape.iter().all(|&dim| dim <= i64::from(i32::MAX)),
        "Flatten input dimensions must fit int32"
    );
    let axis = node.int_attr("axis", 1);
    require!(
        (-i64::from(i32::MAX)..=i64::from(i32::MAX)).contains(&axis),
        "Flatten axis exceeds the native int32 rank range"
    );
    require!(
        node.since_version() >= 11 || axis >= 0,
        "Flatten before schema11 requires nonnegative axis"
    );
    if node.input_rank_known(0) == Some(true) {
        let axis = match checked_flatten_axis(info.shape.len(), axis, node.since_version() >= 11) {
            Ok(axis) => axis,
            Err(error) => deny!("{error}"),
        };
        for dims in [&info.shape[..axis], &info.shape[axis..]] {
            if dims.iter().all(|&dim| dim >= 0) {
                let checked = checked_flatten_product(dims);
                require!(checked.is_ok(), "{}", checked.unwrap_err());
            }
        }
    }
    Ok(())
}

fn flatten_compile_shape_safety(node: &NodeView) -> crate::registry::CompileShapeSafety {
    if node.input_rank_known(0) == Some(true)
        && node
            .input_info(0)
            .is_some_and(|info| info.shape.iter().all(|&dim| dim > 0))
    {
        crate::registry::CompileShapeSafety::Shapeless
    } else {
        crate::registry::CompileShapeSafety::ShapeKeyedOnly {
            reason: "concrete flattened products must not freeze across changing input geometry",
        }
    }
}

fn expand_claim(node: &NodeView) -> ClaimResult {
    require!(
        node.num_inputs() == 2 && node.num_outputs() == 1,
        "Expand expects 2 inputs and 1 output, got {}in/{}out",
        node.num_inputs(),
        node.num_outputs()
    );
    let (data, out) = match movable_io(node) {
        Some(v) => v,
        None => deny!("missing tensor type/shape info on the data input or output"),
    };
    require!(
        is_movable(data),
        "data dtype {} is not a movable dtype supported on GPU",
        crate::registry::ort_dtype_name(data)
    );
    require!(
        out == data,
        "output dtype {} must match data dtype {}",
        crate::registry::ort_dtype_name(out),
        crate::registry::ort_dtype_name(data)
    );
    require!(
        node.input_info(1).map(|i| i.dtype)
            == Some(ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64),
        "Expand `shape` input must be int64 (runtime shape-derived values resolve at trace time; a data-dependent shape falls back to eager)"
    );
    Ok(())
}

fn slice_claim(node: &NodeView) -> ClaimResult {
    let nin = node.num_inputs();
    let legacy = node.since_version() < 10;
    require!(
        (if legacy {
            nin == 1
        } else {
            (3..=5).contains(&nin)
        }) && node.num_outputs() == 1,
        "Slice expects 1 legacy or 3-5 modern inputs and 1 output, got {}in/{}out",
        nin,
        node.num_outputs()
    );
    let (data, out) = match movable_io(node) {
        Some(v) => v,
        None => deny!("missing tensor type/shape info on the data input or output"),
    };
    let data_info = node.input_info(0).ok_or("Slice: missing input metadata")?;
    require!(
        node.input_rank_known(0) != Some(true) || !data_info.shape.is_empty(),
        "Slice: cannot slice a scalar"
    );
    require!(
        data_info.shape.iter().all(|&dim| dim <= i32::MAX as i64),
        "Slice: data dimensions must fit MLX int32 geometry"
    );
    require!(
        is_movable(data),
        "data dtype {} is not a movable dtype supported on GPU",
        crate::registry::ort_dtype_name(data)
    );
    require!(
        out == data,
        "output dtype {} must match data dtype {}",
        crate::registry::ort_dtype_name(out),
        crate::registry::ort_dtype_name(data)
    );
    let (starts, ends, axes, steps) = if legacy {
        let (have_starts, starts) = node.ints_attr("starts");
        let (have_ends, ends) = node.ints_attr("ends");
        require!(
            have_starts && have_ends,
            "Slice: starts and ends attributes are required"
        );
        let (have_axes, axes) = node.ints_attr("axes");
        let axes = if have_axes {
            axes
        } else {
            (0..starts.len() as i64).collect()
        };
        let steps = vec![1; starts.len()];
        (Some(starts), Some(ends), Some(axes), Some(steps))
    } else {
        require!(
            node.input_present(1) && node.input_present(2),
            "Slice: starts and ends are required"
        );
        let index_dtype = node.input_info(1).map(|info| info.dtype);
        for index in 1..nin {
            if !node.input_present(index) {
                continue;
            }
            let info = match node.input_info(index) {
                Some(info) if is_int_index(info.dtype) => info,
                _ => deny!("Slice: input {index} must be an int32/int64 tensor"),
            };
            require!(
                Some(info.dtype) == index_dtype,
                "Slice: control inputs must have the same dtype"
            );
            require!(
                node.input_rank_known(index) != Some(true) || info.shape.len() == 1,
                "Slice: input {index} must have rank one"
            );
        }
        let starts = node.read_const_ints_any(1);
        let ends = node.read_const_ints_any(2);
        let axes = if node.input_present(3) {
            node.read_const_ints_any(3)
        } else {
            starts
                .as_ref()
                .map(|starts| (0..starts.len() as i64).collect())
        };
        let steps = if node.input_present(4) {
            node.read_const_ints_any(4)
        } else {
            starts.as_ref().map(|starts| vec![1; starts.len()])
        };
        if let Some(steps) = &steps {
            require!(!steps.contains(&0), "Slice: steps must not be zero");
        }
        (starts, ends, axes, steps)
    };
    if node.input_rank_known(0) == Some(true)
        && let (Some(starts), Some(ends), Some(axes), Some(steps)) = (starts, ends, axes, steps)
    {
        checked_slice_axes(data_info.shape.len(), &starts, &ends, &axes, &steps)?;
    }
    Ok(())
}

fn split_claim(node: &NodeView) -> ClaimResult {
    let nin = node.num_inputs();
    require!(
        nin != 0 && nin <= 2 && node.num_outputs() != 0,
        "Split expects 1-2 inputs and 1+ outputs, got {}in/{}out",
        nin,
        node.num_outputs()
    );
    let data = match node.input_info(0) {
        Some(d)
            if is_movable(d.dtype)
                && (node.input_rank_known(0) != Some(true) || !d.shape.is_empty()) =>
        {
            d
        }
        Some(d) => deny!(
            "Split: data dtype {} must be a movable GPU dtype and data must be non-scalar (shape {:?})",
            crate::registry::ort_dtype_name(d.dtype),
            d.shape
        ),
        None => deny!("missing data tensor type/shape info"),
    };
    require!(
        data.shape.iter().all(|&dim| dim <= i32::MAX as i64),
        "Split: data dimensions must fit MLX int32 geometry"
    );
    for i in 0..node.num_outputs() {
        match node.output_info(i) {
            Some(o) if o.dtype == data.dtype => {}
            Some(o) => deny!(
                "Split: output[{i}] dtype {} must match data dtype {}",
                crate::registry::ort_dtype_name(o.dtype),
                crate::registry::ort_dtype_name(data.dtype)
            ),
            None => deny!("Split: output[{i}] has no tensor type/shape info"),
        }
    }
    let axis = node.int_attr("axis", 0);
    let axis_size = if node.input_rank_known(0) == Some(true) {
        let rank = data.shape.len() as i64;
        require!(
            axis >= -rank && axis < rank,
            "axis {axis} is out of range for rank {rank}"
        );
        let extent = data.shape[norm_axis(axis, rank as i32) as usize];
        (extent >= 0).then_some(extent)
    } else {
        None
    };

    // Explicit per-section sizes: opset-13 `split` input or opset<13 `split` INTS attribute.
    let (have_sizes, sizes) = if nin == 2 && node.input_present(1) {
        require!(
            node.since_version() >= 13,
            "Split: legacy optional floating split input has no supported length contract"
        );
        let info = node.input_info(1).ok_or("Split: missing split metadata")?;
        require!(
            info.dtype == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64,
            "Split: split sizes must be int64"
        );
        require!(
            node.input_rank_known(1) != Some(true) || info.shape.len() == 1,
            "Split: split sizes must have rank one"
        );
        (true, node.read_const_int64(1))
    } else {
        let (present, s) = node.ints_attr("split");
        (present, present.then_some(s))
    };
    let num_outputs = node
        .has_attr("num_outputs")
        .then(|| node.int_attr("num_outputs", 0));
    if have_sizes && num_outputs.is_some() {
        deny!("Split: split sizes and num_outputs are mutually exclusive");
    }
    if !have_sizes || sizes.is_some() {
        checked_split_sizes(axis_size, node.num_outputs(), sizes.as_deref(), num_outputs)?;
    } else if let Some(info) = node.input_info(1)
        && info.shape.len() == 1
        && info.shape[0] >= 0
    {
        require!(
            info.shape[0] as usize == node.num_outputs(),
            "Split: `split` count {} must equal the number of outputs {}",
            info.shape[0],
            node.num_outputs()
        );
    }
    Ok(())
}

fn split_to_sequence_claim(node: &NodeView) -> ClaimResult {
    require!(
        (1..=2).contains(&node.num_inputs()) && node.num_outputs() == 1,
        "SplitToSequence-24 expects 1-2 inputs and 1 output"
    );
    match node.input_kind(0) {
        Some(ValueKind::Tensor(info)) if is_movable(info.dtype) && !info.shape.is_empty() => {}
        Some(ValueKind::Tensor(info)) => deny!(
            "SplitToSequence-24: input must be a non-scalar movable tensor, got {} shape {:?}",
            crate::registry::ort_dtype_name(info.dtype),
            info.shape
        ),
        Some(kind) => deny!(
            "SplitToSequence-24: input must be a tensor, got {}",
            kind.description()
        ),
        None => deny!("SplitToSequence-24: input is absent"),
    }
    match node.output_kind(0) {
        Some(ValueKind::Sequence(inner)) if inner.tensor().is_some() => deny!(
            "SplitToSequence-24: sequence(tensor) output cannot cross the ORT 1.29 plugin EP boundary; KernelContext_GetOutput only allocates tensors"
        ),
        Some(kind) => deny!(
            "SplitToSequence-24: {} output is unsupported; only sequence(tensor) is valid and container outputs stay on CPU",
            kind.description()
        ),
        None => deny!("SplitToSequence-24: output type information is unavailable"),
    }
}

pub(crate) fn tile_payload_dtype(
    node: &NodeView,
) -> Result<ort::ONNXTensorElementDataType, std::borrow::Cow<'static, str>> {
    validate_tile_payload(movable_io(node), node.since_version())
}

pub(crate) fn validate_tile_payload(
    types: Option<(
        ort::ONNXTensorElementDataType,
        ort::ONNXTensorElementDataType,
    )>,
    since_version: i32,
) -> Result<ort::ONNXTensorElementDataType, std::borrow::Cow<'static, str>> {
    let (data, out) = match types {
        Some(v) => v,
        None => deny!("missing tensor type/shape info on the data input or output"),
    };
    require!(
        is_movable(data),
        "data dtype {} is not a movable dtype supported on GPU",
        crate::registry::ort_dtype_name(data)
    );
    require!(
        out == data,
        "output dtype {} must match data dtype {}",
        crate::registry::ort_dtype_name(out),
        crate::registry::ort_dtype_name(data)
    );
    require!(
        data != ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16
            || since_version >= 13,
        "Tile: bfloat16 requires opset 13"
    );
    Ok(data)
}

fn tile_claim(node: &NodeView) -> ClaimResult {
    require!(
        node.since_version() >= 6,
        "Tile-1 has a distinct unsupported three-input signature"
    );
    require!(
        node.num_inputs() == 2 && node.num_outputs() == 1,
        "Tile expects 2 inputs and 1 output, got {}in/{}out",
        node.num_inputs(),
        node.num_outputs()
    );
    require!(node.num_attributes() == Some(0), "Tile has no attributes");
    let data = tile_payload_dtype(node)?;
    let repeats_info = node.input_info(1).ok_or("Tile: missing repeats metadata")?;
    require!(
        repeats_info.dtype == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64,
        "Tile: repeats must be int64"
    );
    require!(
        node.input_rank_known(1) != Some(true) || repeats_info.shape.len() == 1,
        "Tile: repeats must have rank one"
    );
    let info = node.input_info(0).ok_or("Tile: missing data metadata")?;
    if !node.is_const_int64(1) {
        if node.input_rank_known(0) == Some(true) {
            if node.input_rank_known(1) == Some(true) && repeats_info.shape[0] >= 0 {
                checked_tile_repeats_metadata(&repeats_info.shape, info.shape.len())?;
            }
            if node.output_rank_known(0) == Some(true) {
                require!(
                    node.output_info(0)
                        .ok_or("Tile: missing output metadata")?
                        .shape
                        .len()
                        == info.shape.len(),
                    "Tile: output rank must match data rank"
                );
            }
        }
        if node.input_rank_known(1) == Some(true)
            && repeats_info.shape[0] >= 0
            && node.output_rank_known(0) == Some(true)
        {
            require!(
                node.output_info(0)
                    .ok_or("Tile: missing output metadata")?
                    .shape
                    .len() as i64
                    == repeats_info.shape[0],
                "Tile: output rank must match repeats length"
            );
        }
        return Ok(());
    }
    let repeats = node
        .read_const_int64(1)
        .ok_or("Tile: cannot read constant repeats")?;
    require!(
        node.input_rank_known(1) == Some(true),
        "Tile: repeats must have rank one"
    );
    let shape = if node.input_rank_known(0) == Some(true) {
        info.shape
    } else {
        vec![-1; repeats.len()]
    };
    let output = checked_tile_geometry(
        &shape,
        &repeats_info.shape,
        &repeats,
        mlx_dtype_from_onnx(data),
        true,
    )?;
    if node.output_rank_known(0) == Some(true) {
        let declared = node.output_info(0).ok_or("Tile: missing output metadata")?;
        require!(
            declared.shape.len() == output.len(),
            "Tile: output rank must match data rank"
        );
        require!(
            declared
                .shape
                .iter()
                .zip(&output)
                .all(|(&declared, &actual)| { declared < 0 || actual < 0 || declared == actual }),
            "Tile: declared output dimensions disagree with repeats"
        );
    }
    Ok(())
}

fn tile_compile_shape_safety(node: &NodeView) -> crate::registry::CompileShapeSafety {
    if !node.is_const_int64(1) {
        return crate::registry::CompileShapeSafety::EagerOnly {
            reason: "Tile live repeats require fresh host values on every invocation",
        };
    }
    let shape = node
        .input_info(0)
        .filter(|_| node.input_rank_known(0) == Some(true))
        .map(|info| info.shape);
    tile_geometry_shape_safety(shape.as_deref(), node.read_const_int64(1).as_deref())
}

fn tile_geometry_shape_safety(
    shape: Option<&[i64]>,
    repeats: Option<&[i64]>,
) -> crate::registry::CompileShapeSafety {
    use crate::registry::CompileShapeSafety;
    if shape.is_some_and(|shape| shape.iter().any(|&extent| extent > i64::from(i32::MAX))) {
        return CompileShapeSafety::EagerOnly {
            reason: "Tile metadata-only empty output must not pre-import wide input extents",
        };
    }
    if shape.is_some_and(|shape| shape.iter().all(|&extent| extent > 0))
        && repeats.is_some_and(|repeats| repeats.iter().all(|&repeat| repeat > 0))
    {
        CompileShapeSafety::Shapeless
    } else {
        CompileShapeSafety::ShapeKeyedOnly {
            reason: "conservatively shape-keyed: Tile emits concrete reshape/broadcast and empty output geometry",
        }
    }
}

fn pad_claim(node: &NodeView) -> ClaimResult {
    let nin = node.num_inputs();
    let version = node.since_version();
    let legacy = version < 11;
    let max_inputs = if legacy {
        1
    } else if version < 18 {
        3
    } else {
        4
    };
    require!(
        (if legacy {
            nin == 1
        } else {
            (2..=max_inputs).contains(&nin)
        }) && node.num_outputs() == 1,
        "Pad: wrong input/output counts for schema {version}: {}in/{}out",
        nin,
        node.num_outputs()
    );
    let (data, out) = match movable_io(node) {
        Some(v) => v,
        None => deny!("missing tensor type/shape info on the data input or output"),
    };
    require!(
        is_movable(data),
        "data dtype {} is not a supported numeric or Boolean movement dtype",
        crate::registry::ort_dtype_name(data)
    );
    require!(
        out == data,
        "output dtype {} must match data dtype {}",
        crate::registry::ort_dtype_name(out),
        crate::registry::ort_dtype_name(data)
    );
    let mode = node.string_attr("mode", "constant");
    require!(
        matches!(mode.as_str(), "constant" | "reflect" | "edge" | "wrap"),
        "Pad: mode must be constant, reflect, edge, or wrap"
    );
    require!(
        mode != "wrap" || version >= 19,
        "Pad: wrap mode requires schema 19 or newer"
    );
    let data_info = node.input_info(0).ok_or("Pad: missing data metadata")?;
    let rank = if node.input_rank_known(0) == Some(true) {
        Some(data_info.shape.len())
    } else {
        None
    };
    require!(
        data_info
            .shape
            .iter()
            .all(|&extent| extent <= i64::from(i32::MAX)),
        "Pad: input dimensions must fit native int32 geometry"
    );
    let pads = if legacy {
        let attribute = if version == 1 { "paddings" } else { "pads" };
        let (present, values) = node.ints_attr(attribute);
        require!(present, "Pad: missing {attribute} attribute");
        Some(values)
    } else {
        let info = node.input_info(1).ok_or("Pad: missing pads metadata")?;
        require!(
            info.dtype == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64
                && (node.input_rank_known(1) != Some(true) || info.shape.len() == 1),
            "Pad: pads must be a rank-one int64 tensor"
        );
        node.read_const_int64(1)
    };
    let known_count = |index: usize| {
        node.input_info(index).and_then(|info| {
            if node.input_rank_known(index) == Some(true) && info.shape.len() == 1 {
                usize::try_from(info.shape[0]).ok()
            } else {
                None
            }
        })
    };
    let pad_count = pads.as_ref().map(Vec::len).or_else(|| known_count(1));
    if let Some(count) = pad_count {
        require!(count % 2 == 0, "Pad: pads length must be even");
    }
    let axes = if node.input_present(3) {
        let info = node.input_info(3).ok_or("Pad: missing axes metadata")?;
        require!(
            is_int_index(info.dtype)
                && (node.input_rank_known(3) != Some(true) || info.shape.len() == 1),
            "Pad: axes must be a rank-one int32 or int64 tensor"
        );
        node.read_const_ints_any(3)
    } else {
        rank.map(|rank| (0..rank as i64).collect())
    };
    let axis_count = axes.as_ref().map(Vec::len).or_else(|| known_count(3));
    if let (Some(pad_count), Some(axis_count)) = (pad_count, axis_count) {
        require!(
            axis_count.checked_mul(2) == Some(pad_count),
            "Pad: pads length must be twice the axes length"
        );
    }
    if let (Some(rank), Some(axes)) = (rank, axes.as_ref()) {
        let count = pad_count.unwrap_or(
            axes.len()
                .checked_mul(2)
                .ok_or("Pad: axes count overflows")?,
        );
        checked_pad_axes(rank, axes, count)?;
        if let Some(pads) = pads.as_ref()
            && data_info.shape.iter().all(|&extent| extent >= 0)
        {
            checked_pad_geometry(&data_info.shape, axes, pads, &mode)?;
        }
    }
    if nin >= 3 && node.input_present(2) {
        match node.input_info(2) {
            Some(cv)
                if cv.dtype == data
                    && (node.input_rank_known(2) != Some(true) || cv.shape.is_empty()) => {}
            Some(cv) => deny!(
                "Pad: constant_value must be a scalar with dtype {} (got {} shape {:?})",
                crate::registry::ort_dtype_name(data),
                crate::registry::ort_dtype_name(cv.dtype),
                cv.shape
            ),
            None => deny!("Pad: constant_value input has no tensor type/shape info"),
        }
    }
    Ok(())
}

fn identity_claim(node: &NodeView) -> ClaimResult {
    require!(
        node.num_inputs() == 1 && node.num_outputs() == 1,
        "Identity expects 1 input and 1 output"
    );
    match (node.input_kind(0), node.output_kind(0)) {
        (Some(ValueKind::Tensor(_)), Some(ValueKind::Tensor(_))) => {}
        (Some(input), Some(output)) => deny!(
            "Identity: {} -> {} crosses the ORT 1.29 plugin EP boundary; sequence/optional outputs cannot be bound by KernelContext_GetOutput",
            input.description(),
            output.description()
        ),
        _ => deny!("Identity: missing input or output type information"),
    }
    movable_view_claim(node)
}

fn range_claim(node: &NodeView) -> ClaimResult {
    require!(
        node.num_inputs() == 3 && node.num_outputs() == 1,
        "Range expects 3 inputs and 1 output, got {}in/{}out",
        node.num_inputs(),
        node.num_outputs()
    );
    let ty = match node.input_info(0) {
        Some(i) if is_range_type(i.dtype) => i.dtype,
        Some(i) => deny!(
            "Range: start dtype {} unsupported",
            crate::registry::ort_dtype_name(i.dtype)
        ),
        None => deny!("missing `start` tensor type/shape info"),
    };
    let out = match node.output_info(0) {
        Some(o) if o.dtype == ty && o.shape.len() == 1 => o,
        Some(o) => deny!(
            "Range: output dtype {} must equal start dtype {} and output must be 1-D (shape {:?})",
            crate::registry::ort_dtype_name(o.dtype),
            crate::registry::ort_dtype_name(ty),
            o.shape
        ),
        None => deny!("missing output tensor type/shape info"),
    };
    for i in 0..3 {
        let info = node.input_info(i).ok_or("Range: missing scalar metadata")?;
        require!(
            info.dtype == ty && node.input_rank_known(i) == Some(true) && info.shape.is_empty(),
            "Range: inputs must be scalars with the same dtype"
        );
    }
    if ty == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64 {
        let mut bounds = [0i64; 3];
        for (i, value) in bounds.iter_mut().enumerate() {
            let values = match node.read_const_int64(i) {
                Some(values) if values.len() == 1 => values,
                _ => return Ok(()), // Live INT64 bounds are evaluated at an eager boundary.
            };
            *value = values[0];
        }
        let count = integer_range_count(bounds[0], bounds[1], bounds[2])?;
        require!(
            out.shape[0] == count as i64,
            "Range: inferred output extent {} disagrees with exact integer count {count}",
            out.shape[0]
        );
        return Ok(());
    }
    // Range's output extent is value-dependent, so the bounds MUST be constant scalar initializers:
    // claiming a runtime-bounded Range puts it in the compiled closure where its output feeds a
    // downstream shape/axes reader (e.g. ReduceMean axes in expanded RMSNormalization), forcing a
    // mid-compile-trace eval that aborts `mlx_compile`. Runtime bounds stay on CPU (vision-only value).
    let (start, limit, delta) = match (
        node.read_const_scalar_f64(0),
        node.read_const_scalar_f64(1),
        node.read_const_scalar_f64(2),
    ) {
        (Some(a), Some(b), Some(c)) => (a, b, c),
        _ => deny!(
            "Range: `start`/`limit`/`delta` must be constant scalar initializers; this node's are runtime values — stays on CPU"
        ),
    };
    require!(delta != 0.0, "Range: `delta` must be non-zero");
    if node.since_version() >= 27
        && (ty == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16
            || ty == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16)
    {
        let stash = node.int_attr("stash_type", 1);
        require!(
            stash == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT as i64
                || stash
                    == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE as i64,
            "Range-27 fp16/bf16 stash_type must be FLOAT or DOUBLE"
        );
    }
    let count = ((limit - start) / delta).ceil().max(0.0);
    require!(
        count.is_finite() && count <= i32::MAX as f64,
        "Range: computed element count {count} is not finite or exceeds i32::MAX"
    );
    require!(
        out.shape[0] == count as i64,
        "Range: output extent {} must equal computed element count {}",
        out.shape[0],
        count as i64
    );
    Ok(())
}

fn shape_claim(node: &NodeView) -> ClaimResult {
    require!(
        node.num_inputs() == 1 && node.num_outputs() == 1,
        "Shape expects 1 input and 1 output, got {}in/{}out",
        node.num_inputs(),
        node.num_outputs()
    );
    let data = match node.input_info(0) {
        Some(d) if d.dtype != 0 => d,
        Some(_) => deny!("Shape: data tensor dtype must be known"),
        None => deny!("missing data tensor type/shape info"),
    };
    let out = match node.output_info(0) {
        Some(o)
            if o.dtype == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64
                && o.shape.len() == 1 =>
        {
            o
        }
        Some(o) => deny!(
            "Shape: output must be a 1-D int64 tensor (got {} shape {:?})",
            crate::registry::ort_dtype_name(o.dtype),
            o.shape
        ),
        None => deny!("missing output tensor type/shape info"),
    };
    let rank = data.shape.len() as i64;
    let (s, e) = shape_interval(rank, node.int_attr("start", 0), node.int_attr("end", rank));
    require!(
        node.input_rank_known(0) != Some(true) || out.shape[0] < 0 || out.shape[0] == e - s,
        "Shape: output extent {} must equal the sliced rank interval {} (start..end)",
        out.shape[0],
        e - s
    );
    Ok(())
}

fn size_claim(node: &NodeView) -> ClaimResult {
    require!(
        node.num_inputs() == 1 && node.num_outputs() == 1,
        "Size expects 1 input and 1 output, got {}in/{}out",
        node.num_inputs(),
        node.num_outputs()
    );
    require!(
        matches!(node.input_info(0), Some(d) if d.dtype != 0),
        "Size: data must have known tensor type information"
    );
    require!(
        matches!(node.output_info(0), Some(o)
            if o.dtype == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64
                && o.shape.is_empty()),
        "Size: output must be a scalar int64 tensor"
    );
    Ok(())
}

fn metadata_compile_shape_safety(node: &NodeView) -> crate::registry::CompileShapeSafety {
    use crate::registry::CompileShapeSafety;
    match node.input_info(0) {
        Some(info)
            if !is_movable(info.dtype) || info.shape.iter().any(|&dim| dim > i32::MAX as i64) =>
        {
            CompileShapeSafety::EagerOnly {
                reason: "conservatively reads non-MLX tensor geometry without converting payloads",
            }
        }
        Some(info)
            if node.input_rank_known(0) == Some(true)
                && info.shape.iter().all(|&dim| dim >= 0)
                && node.input_is_host_readable(0) =>
        {
            CompileShapeSafety::Shapeless
        }
        Some(_) => CompileShapeSafety::ShapeKeyedOnly {
            reason: "runtime and intermediate tensor metadata must not freeze in shapeless traces",
        },
        None => CompileShapeSafety::EagerOnly {
            reason: "conservatively requires live tensor metadata",
        },
    }
}

fn constant_of_shape_claim(node: &NodeView) -> ClaimResult {
    require!(
        node.num_inputs() == 1 && node.num_outputs() == 1,
        "ConstantOfShape expects 1 input and 1 output, got {}in/{}out",
        node.num_inputs(),
        node.num_outputs()
    );
    let out = match node.output_info(0) {
        Some(o) if is_movable(o.dtype) => o.dtype,
        Some(o) => deny!(
            "output dtype {} is not a movable dtype supported on GPU",
            crate::registry::ort_dtype_name(o.dtype)
        ),
        None => deny!("missing output tensor type/shape info"),
    };
    // A constant `input` is checked here; a runtime one is resolved at trace time by `read_ints`,
    // exactly as Reshape's `shape` input is. Refusing the runtime form is what used to strand a
    // zero-init state tensor on CPU, and each such node is a partition boundary: on a dynamic-shape
    // Mamba model that split the graph into 79 fused subgraphs instead of 1.
    match node.read_const_int64(0) {
        Some(shape) => require!(
            shape.iter().all(|&d| d >= 0 && d <= i32::MAX as i64),
            "ConstantOfShape: shape dims must be in [0, i32::MAX] (got {:?})",
            shape
        ),
        None => require!(
            node.input_info(0).map(|i| i.dtype)
                == Some(ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64),
            "ConstantOfShape: `input` must be int64 (a runtime shape-derived value resolves at trace time)"
        ),
    }
    // ONNX defines `value` as a one-element tensor giving both the fill value and the output dtype;
    // omitting it means fp32 zero. Both forms are handled, so long as the output dtype is one MLX
    // can carry. `is_movable` rather than `is_mlx_supported`: this op only fills, never computes, so
    // it takes the same dtype set as the data-movement ops — float64 included, which GetCapability
    // then isolates onto an MLX CPU stream.
    require!(
        crate::registry::is_movable(out),
        "ConstantOfShape: output dtype {} is not an MLX dtype",
        crate::registry::ort_dtype_name(out)
    );
    // A `value` this cannot rebuild byte-exactly in an MLX dtype is left to CPU rather than
    // silently narrowed.
    if let Some(v) = node.attr_tensor_dtype("value") {
        require!(
            crate::registry::is_movable(v),
            "ConstantOfShape: `value` dtype {} is not an MLX dtype",
            crate::registry::ort_dtype_name(v)
        );
    }
    Ok(())
}

fn constant_of_shape_compile_shape_safety(node: &NodeView) -> crate::registry::CompileShapeSafety {
    if node.is_constant_initializer(0) {
        crate::registry::CompileShapeSafety::Shapeless
    } else {
        crate::registry::CompileShapeSafety::ShapeKeyedOnly {
            reason: "runtime fill shape requires shape/value guards, not shapeless replay",
        }
    }
}

// The per-axis claim guards below index parallel arrays (input.shape/output.shape/sizes/scales) by
// axis and use the axis index in their decline messages — an iterator form would be less clear.
#[allow(clippy::needless_range_loop)]
fn resize_claim(node: &NodeView) -> ClaimResult {
    require!(
        node.num_inputs() >= 3 && node.num_inputs() <= 4 && node.num_outputs() == 1,
        "Resize expects 3-4 inputs and 1 output, got {}in/{}out",
        node.num_inputs(),
        node.num_outputs()
    );
    let (input, output) = match (node.input_info(0), node.output_info(0)) {
        (Some(input), Some(output)) => (input, output),
        _ => deny!("missing tensor type/shape info on the input or output"),
    };
    let rank = input.shape.len();
    // Spatial dims MAY be dynamic (<=0): the concrete input size is resolved at trace time from
    // `ctx.shape_of`, so we only require the resize TARGET to be statically determinable (constant
    // `scales`/`sizes`) and the batch/channel dims to be untouched. Rank must be known and match.
    require!(
        is_mlx_float(input.dtype),
        "Resize: input dtype {} must be an MLX float (fp32/fp16/bf16)",
        crate::registry::ort_dtype_name(input.dtype)
    );
    require!(
        output.dtype == input.dtype,
        "Resize: output dtype {} must match input dtype {}",
        crate::registry::ort_dtype_name(output.dtype),
        crate::registry::ort_dtype_name(input.dtype)
    );
    require!(
        rank != 0 && rank <= 4,
        "Resize: input rank {rank} unsupported (only 1-D..4-D tensors are claimed)"
    );
    require!(
        output.shape.len() == rank,
        "Resize: output rank {} must equal input rank {rank}",
        output.shape.len()
    );
    let mode = node.string_attr("mode", "nearest");
    require!(
        mode == "nearest" || mode == "linear",
        "Resize: only mode=nearest|linear are claimed (got mode={mode}; cubic stays on CPU)"
    );
    let coordinate_mode = node.string_attr("coordinate_transformation_mode", "half_pixel");
    require!(
        matches!(
            coordinate_mode.as_str(),
            "half_pixel" | "asymmetric" | "align_corners" | "pytorch_half_pixel"
        ),
        "Resize: coordinate_transformation_mode={coordinate_mode} is unsupported (only half_pixel|asymmetric|align_corners|pytorch_half_pixel are claimed)"
    );
    require!(
        mode != "nearest"
            || matches!(
                node.string_attr("nearest_mode", "round_prefer_floor")
                    .as_str(),
                "round_prefer_floor" | "round_prefer_ceil" | "floor" | "ceil"
            ),
        "Resize: nearest_mode={} is unsupported (only round_prefer_floor|round_prefer_ceil|floor|ceil are claimed)",
        node.string_attr("nearest_mode", "round_prefer_floor")
    );
    require!(
        node.int_attr("exclude_outside", 0) == 0,
        "Resize: exclude_outside=1 is unsupported — stays on CPU"
    );
    require!(
        node.int_attr("antialias", 0) == 0,
        "Resize: antialias=1 is unsupported — stays on CPU"
    );
    require!(
        !node.has_attr("axes"),
        "Resize: the `axes` attribute is unsupported (full-rank scales/sizes only) — stays on CPU"
    );
    require!(
        node.string_attr("keep_aspect_ratio_policy", "stretch") == "stretch",
        "Resize: only keep_aspect_ratio_policy=stretch is claimed — stays on CPU"
    );
    require!(
        !node.input_present(1),
        "Resize: the `roi` input (tf_crop_and_resize) is unsupported — stays on CPU"
    );
    let has_scales = node.input_present(2);
    let has_sizes = node.input_present(3);
    require!(
        has_scales != has_sizes,
        "Resize: exactly one of `scales` or `sizes` must be provided (not both/neither)"
    );
    // For rank>=3 the batch/channel axes (0,1) must NOT be resized — MLX resize here is
    // spatial-only, and a batch/channel target can't be verified against a dynamic input dim.
    let bc = if rank >= 3 { 2 } else { 0 };
    if has_sizes {
        // Constant `sizes` give the exact (static) output extents directly.
        let sizes = match node.read_const_int64(3) {
            Some(s) if s.len() == rank && s.iter().all(|&v| v >= 1) => s,
            Some(s) => deny!(
                "Resize: `sizes` must have rank {rank} with all entries >= 1 (got {:?})",
                s
            ),
            None => deny!(
                "Resize: only constant (initializer) `scales`/`sizes` are claimed; this node's are runtime/dynamic — precompute them or export a static-shape model"
            ),
        };
        for ax in 0..bc {
            // Batch/channel: input dim must be static and unchanged.
            require!(
                input.shape[ax] > 0 && sizes[ax] == input.shape[ax],
                "Resize: batch/channel dim {ax} must be static and unchanged (input {}, requested size {})",
                input.shape[ax],
                sizes[ax]
            );
        }
        for ax in bc..rank {
            // Where the output spatial dim is statically known, it must equal the requested size.
            require!(
                output.shape[ax] <= 0 || output.shape[ax] == sizes[ax],
                "Resize: static output spatial dim {ax} ({}) must equal requested size {}",
                output.shape[ax],
                sizes[ax]
            );
        }
        Ok(())
    } else {
        // Constant `scales`: the concrete output is `floor(scale * input)` computed at trace time.
        let scales = match node.read_const_f32(2) {
            Some(s) if s.len() == rank && s.iter().all(|&v| v > 0.0) => s,
            Some(s) => deny!(
                "Resize: `scales` must have rank {rank} with all entries > 0 (got {:?})",
                s
            ),
            None => deny!(
                "Resize: only constant (initializer) `scales`/`sizes` are claimed; this node's are runtime/dynamic — precompute them or export a static-shape model"
            ),
        };
        for ax in 0..bc {
            // Batch/channel scale must be exactly 1 (no resize).
            require!(
                (scales[ax] - 1.0).abs() <= f32::EPSILON,
                "Resize: batch/channel dim {ax} must not be resized (scale {} != 1)",
                scales[ax]
            );
        }
        for ax in bc..rank {
            // Where both input and output spatial dims are static, verify the computed extent.
            if input.shape[ax] > 0 && output.shape[ax] > 0 {
                let computed = (scales[ax] as f64 * input.shape[ax] as f64).floor() as i64;
                require!(
                    computed == output.shape[ax],
                    "Resize: spatial dim {ax} computed extent {computed} (floor(scale*input)) must equal static output {}",
                    output.shape[ax]
                );
            }
        }
        Ok(())
    }
}

fn upsample_claim(node: &NodeView) -> ClaimResult {
    require!(
        node.num_inputs() == 2 && node.num_outputs() == 1,
        "Upsample-9 expects 2 inputs and 1 output, got {}in/{}out",
        node.num_inputs(),
        node.num_outputs()
    );
    let (input, scales_info, output) =
        match (node.input_info(0), node.input_info(1), node.output_info(0)) {
            (Some(input), Some(scales), Some(output)) => (input, scales, output),
            _ => deny!("Upsample: missing tensor type/shape info"),
        };
    let rank = input.shape.len();
    require!(
        rank != 0
            && rank <= 4
            && input.shape.iter().all(|&d| d > 0)
            && output.shape.iter().all(|&d| d > 0),
        "Upsample: input/output must have fully-static positive rank-1..4 shapes"
    );
    require!(
        output.shape.len() == rank && output.dtype == input.dtype,
        "Upsample: output rank and dtype must match the input"
    );
    let mode = node.string_attr("mode", "nearest");
    require!(
        mode == "nearest" || mode == "linear",
        "Upsample: only legacy mode=nearest|linear is supported"
    );
    require!(
        (mode == "nearest" && is_movable(input.dtype))
            || (mode == "linear" && is_mlx_float(input.dtype)),
        "Upsample: nearest supports movable numeric/bool dtypes; linear requires fp32/fp16/bf16"
    );
    require!(
        scales_info.dtype == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT
            && scales_info.shape == vec![rank as i64],
        "Upsample: scales must be a rank-sized float32 vector"
    );
    let scales = match node.read_const_f32(1) {
        Some(scales) if scales.len() == rank && scales.iter().all(|&scale| scale > 0.0) => scales,
        Some(scales) => deny!(
            "Upsample: scales must contain {rank} positive values (got {:?})",
            scales
        ),
        None => deny!(
            "Upsample: scales must be a constant float32 initializer because output shape is scale-dependent"
        ),
    };
    for (axis, &scale) in scales.iter().enumerate().take(rank) {
        let expected = (input.shape[axis] as f64 * scale as f64).floor() as i64;
        require!(
            output.shape[axis] == expected,
            "Upsample: output dim {axis} ({}) must equal floor(input {} * scale {}) = {expected}",
            output.shape[axis],
            input.shape[axis],
            scale
        );
    }
    Ok(())
}

fn where_claim(node: &NodeView) -> ClaimResult {
    require!(
        node.num_inputs() == 3 && node.num_outputs() == 1,
        "Where expects 3 inputs and 1 output, got {}in/{}out",
        node.num_inputs(),
        node.num_outputs()
    );
    let (cond, x, y, out) = match (
        node.input_info(0),
        node.input_info(1),
        node.input_info(2),
        node.output_info(0),
    ) {
        (Some(a), Some(b), Some(c), Some(d)) => (a, b, c, d),
        _ => deny!("missing tensor type/shape info on an input or the output"),
    };
    let is_bool = |t| t == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL;
    require!(
        is_bool(cond.dtype),
        "Where: `condition` must be bool (got {})",
        crate::registry::ort_dtype_name(cond.dtype)
    );
    require!(
        x.dtype == y.dtype && y.dtype == out.dtype,
        "Where: X/Y/output must share one dtype (got {}, {} -> {})",
        crate::registry::ort_dtype_name(x.dtype),
        crate::registry::ort_dtype_name(y.dtype),
        crate::registry::ort_dtype_name(out.dtype)
    );
    require!(
        is_mlx_numeric(x.dtype) || is_bool(x.dtype),
        "Where: X/Y dtype {} must be numeric or bool",
        crate::registry::ort_dtype_name(x.dtype)
    );
    require!(
        crate::registry::scalar_or_suffix_broadcast(&cond.shape, &out.shape)
            && crate::registry::scalar_or_suffix_broadcast(&x.shape, &out.shape)
            && crate::registry::scalar_or_suffix_broadcast(&y.shape, &out.shape),
        "Where: only scalar or trailing-suffix broadcast to the output shape is supported (cond {:?}, X {:?}, Y {:?} -> {:?})",
        cond.shape,
        x.shape,
        y.shape,
        out.shape
    );
    Ok(())
}

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

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

pub fn register(registry: &mut OpRegistry) {
    shapeless(registry, "Gather", gather_op, gather_like_claim);
    shape_keyed(
        registry,
        "GatherND",
        gathernd_op,
        gathernd_claim,
        "conservatively shape-keyed: GatherND flattening strides and batch offsets depend on concrete input extents",
    );
    shapeless(
        registry,
        "GatherElements",
        gather_elements_op,
        gather_like_claim,
    );
    shape_keyed(
        registry,
        "ScatterElements",
        scatter_elements_op,
        scatter_elements_claim,
        crate::registry::MLX_SCATTER_SHAPE_REASON,
    );
    shape_keyed(
        registry,
        "CenterCropPad",
        center_crop_pad_op,
        center_crop_pad_claim,
        "emits MLX Slice and Pad, which lack Primitive::output_shapes in MLX 0.32.2",
    );
    shapeless(registry, "Compress", compress_op, compress_claim);
    shapeless(
        registry,
        "DepthToSpace",
        depth_to_space_op,
        depth_to_space_claim,
    );
    shapeless(registry, "EyeLike", eye_like_op, eye_like_claim);
    shapeless(
        registry,
        "ReverseSequence",
        reverse_sequence_op,
        reverse_sequence_claim,
    );
    registry.register_shape_classified(
        OpRegistration {
            domain: "",
            op_type: "ScatterND",
            min_opset: K_ANY_OPSET,
            max_opset: K_ANY_OPSET,
            handler: scatter_nd_op,
            claim: scatter_nd_claim,
        },
        |node| scatter_nd_compile_policy(&node.string_attr("reduction", "none")),
    );
    shapeless(
        registry,
        "SpaceToDepth",
        space_to_depth_op,
        space_to_depth_claim,
    );
    shapeless(registry, "Concat", concat_op, concat_claim);
    shapeless(registry, "Reshape", reshape_op, reshape_claim);
    registry.register_shape_classified(
        OpRegistration {
            domain: "",
            op_type: "Transpose",
            min_opset: 1,
            max_opset: K_ANY_OPSET,
            handler: transpose_op,
            claim: transpose_claim,
        },
        transpose_compile_shape_safety,
    );
    for op_type in ["Unsqueeze", "Squeeze"] {
        registry.register_shape_classified(
            OpRegistration {
                domain: "",
                op_type,
                min_opset: 1,
                max_opset: K_ANY_OPSET,
                handler: rank_change_op,
                claim: rank_change_claim,
            },
            rank_change_compile_shape_safety,
        );
    }
    registry.register_shape_classified(
        OpRegistration {
            domain: "",
            op_type: "Flatten",
            min_opset: 1,
            max_opset: K_ANY_OPSET,
            handler: flatten_op,
            claim: flatten_claim,
        },
        flatten_compile_shape_safety,
    );
    shapeless(registry, "Expand", expand_op, expand_claim);
    shape_keyed(
        registry,
        "Slice",
        slice_op,
        slice_claim,
        crate::registry::MLX_SLICE_SHAPE_REASON,
    );
    shape_keyed(
        registry,
        "Split",
        split_op,
        split_claim,
        crate::registry::MLX_SPLIT_SHAPE_REASON,
    );
    registry.register_shapeless(OpRegistration {
        domain: "",
        op_type: "SplitToSequence",
        min_opset: 24,
        max_opset: K_ANY_OPSET,
        handler: split_to_sequence_op,
        claim: split_to_sequence_claim,
    });
    registry.register_shape_classified(
        OpRegistration {
            domain: "",
            op_type: "Tile",
            min_opset: 6,
            max_opset: K_ANY_OPSET,
            handler: tile_op,
            claim: tile_claim,
        },
        tile_compile_shape_safety,
    );
    shape_keyed(
        registry,
        "Pad",
        pad_op,
        pad_claim,
        crate::registry::MLX_PAD_SHAPE_REASON,
    );
    shapeless(registry, "Identity", identity_op, identity_claim);
    shapeless(registry, "Range", range_op, range_claim);
    registry.register_shape_classified(
        OpRegistration {
            domain: "",
            op_type: "Shape",
            min_opset: 1,
            max_opset: K_ANY_OPSET,
            handler: shape_op,
            claim: shape_claim,
        },
        metadata_compile_shape_safety,
    );
    registry.register_shape_classified(
        OpRegistration {
            domain: "",
            op_type: "Size",
            min_opset: 1,
            max_opset: K_ANY_OPSET,
            handler: size_op,
            claim: size_claim,
        },
        metadata_compile_shape_safety,
    );
    registry.register_shape_classified(
        OpRegistration {
            domain: "",
            op_type: "ConstantOfShape",
            min_opset: K_ANY_OPSET,
            max_opset: K_ANY_OPSET,
            handler: constant_of_shape_op,
            claim: constant_of_shape_claim,
        },
        constant_of_shape_compile_shape_safety,
    );
    shapeless(registry, "Where", where_handler, where_claim);
    shapeless(registry, "Resize", resize_op, resize_claim);
    registry.register_shapeless(OpRegistration {
        domain: "",
        op_type: "Upsample",
        min_opset: 9,
        max_opset: 9,
        handler: resize_op,
        claim: upsample_claim,
    });
}

#[cfg(test)]
mod tests {
    #[test]
    fn scatter_add_checks_empty_scalar_and_invalid_geometry() {
        assert!(matches!(
            super::scatter_nd_compile_policy("none"),
            crate::registry::CompileShapeSafety::ShapeKeyedOnly { .. }
        ));
        assert!(matches!(
            super::scatter_nd_compile_policy("add"),
            crate::registry::CompileShapeSafety::EagerOnly { .. }
        ));
        assert_eq!(
            super::scatter_nd_geometry(&[2, 3, 4], &[0, 2], &[0, 4]).unwrap(),
            (2, 6, 4, 0)
        );
        assert_eq!(
            super::scatter_nd_geometry(&[2, 3], &[2], &[]).unwrap(),
            (2, 6, 1, 1)
        );
        assert!(super::scatter_nd_geometry(&[2], &[1, 2], &[1]).is_err());
        assert!(super::scatter_nd_geometry(&[2], &[1, 0], &[1, 2]).is_err());
        assert!(super::scatter_nd_geometry(&[2, 3], &[2, 1], &[2, 2]).is_err());
        assert!(super::scatter_nd_geometry(&[i32::MAX, 2], &[1, 2], &[1]).is_err());
    }

    use super::{
        checked_flatten_axis, checked_flatten_shape, checked_pad_axes, checked_pad_geometry,
        checked_rank_change_axes, checked_split_sizes, checked_tile_geometry,
        checked_tile_repeats_metadata, checked_transpose_perm, integer_range_count,
        tile_geometry_shape_safety,
    };
    use crate::registry::CompileShapeSafety;
    use crate::sys::mlx;

    #[test]
    fn tile_live_metadata_is_validated_before_reading_values() {
        for rank in [0, 1, 2, 65] {
            assert!(checked_tile_repeats_metadata(&[rank as i64], rank).is_ok());
            for shape in [
                vec![],
                vec![1, rank as i64],
                vec![-1],
                vec![rank as i64 + 1],
            ] {
                assert!(checked_tile_repeats_metadata(&shape, rank).is_err());
            }
        }
    }

    #[test]
    fn tile_exact_checked_geometry() {
        let dtype = mlx::mlx_dtype__MLX_FLOAT32;
        for (shape, repeats, expected) in [
            (vec![], vec![], vec![]),
            (vec![2, 3], vec![3, 2], vec![6, 6]),
            (vec![0, 2], vec![i64::MAX, 3], vec![0, 6]),
            (vec![0, i64::MAX], vec![1, 0], vec![0, 0]),
            (vec![i64::MAX, 0], vec![0, i64::MAX], vec![0, 0]),
            (vec![65536, 65536, 0], vec![1, 1, 1], vec![65536, 65536, 0]),
            (
                vec![i64::from(i32::MAX)],
                vec![1],
                vec![i64::from(i32::MAX)],
            ),
        ] {
            assert_eq!(
                checked_tile_geometry(&shape, &[repeats.len() as i64], &repeats, dtype, false)
                    .unwrap(),
                expected
            );
        }
    }

    #[test]
    fn tile_rejects_each_invalid_axis_before_empty_shortcut() {
        let dtype = mlx::mlx_dtype__MLX_UINT64;
        for (shape, repeats, message) in [
            (vec![0], vec![-1], "repeats must be nonnegative"),
            (vec![0, 2], vec![1, i64::MIN], "repeats must be nonnegative"),
            (
                vec![1],
                vec![(1i64 << 32) + 1],
                "output extent exceeds int32",
            ),
            (vec![1], vec![i64::MAX], "output extent exceeds int32"),
            (vec![2], vec![i64::MAX], "output extent overflows int64"),
            (
                vec![0, 2],
                vec![1, i64::MAX],
                "output extent overflows int64",
            ),
            (
                vec![0, 1],
                vec![1, i64::from(i32::MAX) + 1],
                "output extent exceeds int32",
            ),
            (vec![0, i64::MAX], vec![1, 1], "output extent exceeds int32"),
            (
                vec![65536, 65536],
                vec![1, 1],
                "element count exceeds int32 indexing",
            ),
            (vec![-1], vec![0], "negative runtime extent"),
        ] {
            let error =
                checked_tile_geometry(&shape, &[repeats.len() as i64], &repeats, dtype, false)
                    .unwrap_err();
            assert!(error.contains(message), "{error}");
        }
    }

    #[test]
    fn tile_rank_cardinality_and_symbolic_geometry() {
        let dtype = mlx::mlx_dtype__MLX_BOOL;
        for (shape, repeats_shape, repeats, message) in [
            (vec![2], vec![], vec![2], "repeats must have rank one"),
            (vec![2], vec![1, 1], vec![2], "repeats must have rank one"),
            (
                vec![2, 3],
                vec![1],
                vec![2],
                "repeats length must equal data rank",
            ),
            (
                vec![],
                vec![1],
                vec![],
                "repeats length must equal data rank",
            ),
            (
                vec![2],
                vec![0],
                vec![2],
                "repeats length must equal data rank",
            ),
        ] {
            assert!(
                checked_tile_geometry(&shape, &repeats_shape, &repeats, dtype, false)
                    .unwrap_err()
                    .contains(message)
            );
        }
        assert_eq!(
            checked_tile_geometry(&[-1, 2], &[2], &[3, 2], dtype, true).unwrap(),
            [-1, 4]
        );
        assert_eq!(
            checked_tile_geometry(&[-1, 2], &[2], &[0, 2], dtype, true).unwrap(),
            [0, 4]
        );
        assert!(checked_tile_geometry(&[-1], &[1], &[-1], dtype, true).is_err());
    }

    #[test]
    fn tile_capture_is_geometry_specific() {
        assert_eq!(
            tile_geometry_shape_safety(Some(&[]), Some(&[])),
            CompileShapeSafety::Shapeless
        );
        assert_eq!(
            tile_geometry_shape_safety(Some(&[2, 3]), Some(&[2, 1])),
            CompileShapeSafety::Shapeless
        );
        for (shape, repeats) in [
            (None, Some(&[2, 1][..])),
            (Some(&[-1, 3][..]), Some(&[2, 1][..])),
            (Some(&[0, 3][..]), Some(&[2, 1][..])),
            (Some(&[2, 3][..]), Some(&[0, 1][..])),
        ] {
            assert!(matches!(
                tile_geometry_shape_safety(shape, repeats),
                CompileShapeSafety::ShapeKeyedOnly { .. }
            ));
        }
        assert!(matches!(
            tile_geometry_shape_safety(Some(&[0, i64::MAX]), Some(&[1, 0])),
            CompileShapeSafety::EagerOnly { .. }
        ));
    }

    #[test]
    fn split_sizes_validate_without_known_axis_geometry() {
        assert_eq!(
            checked_split_sizes(None, 3, Some(&[2, 0, 4]), None).unwrap(),
            Some(vec![2, 0, 4])
        );
        assert_eq!(checked_split_sizes(None, 2, None, None).unwrap(), None);
        assert!(checked_split_sizes(None, 2, Some(&[i64::MAX, i64::MAX]), None).is_err());
        assert!(checked_split_sizes(None, 2, Some(&[-1, 7]), None).is_err());
        assert!(checked_split_sizes(None, 2, Some(&[]), None).is_err());
        assert!(checked_split_sizes(None, 2, Some(&[2, 4]), Some(2)).is_err());
        for num_outputs in [i64::MIN, -1, 0, 1, 3, i64::MAX] {
            assert!(checked_split_sizes(None, 2, None, Some(num_outputs)).is_err());
        }
    }

    #[test]
    fn transpose_permutations_preserve_rank_and_full_width_bounds() {
        assert_eq!(checked_transpose_perm(0, None).unwrap(), Vec::<i32>::new());
        assert_eq!(
            checked_transpose_perm(0, Some(&[])).unwrap(),
            Vec::<i32>::new()
        );
        assert_eq!(checked_transpose_perm(3, None).unwrap(), vec![2, 1, 0]);
        assert_eq!(
            checked_transpose_perm(3, Some(&[1, 2, 0])).unwrap(),
            vec![1, 2, 0]
        );
        for perm in [
            vec![],
            vec![0],
            vec![0, 0],
            vec![-1, 0],
            vec![0, 2],
            vec![i64::MIN, 0],
            vec![i64::MAX, 0],
            vec![1 << 32, 1],
        ] {
            assert!(checked_transpose_perm(2, Some(&perm)).is_err());
        }
    }

    #[test]
    fn split_sizes_preserve_empty_and_smaller_final_sections() {
        for (extent, count, num_outputs, expected) in [
            (0, 2, None, vec![0, 0]),
            (0, 2, Some(2), vec![0, 0]),
            (1, 2, Some(2), vec![1, 0]),
            (6, 2, None, vec![3, 3]),
            (5, 2, Some(2), vec![3, 2]),
            (4, 3, Some(3), vec![2, 2, 0]),
            (7, 4, Some(4), vec![2, 2, 2, 1]),
        ] {
            assert_eq!(
                checked_split_sizes(Some(extent), count, None, num_outputs).unwrap(),
                Some(expected)
            );
        }
        assert!(checked_split_sizes(Some(5), 2, None, None).is_err());
        assert!(checked_split_sizes(Some(5), 4, None, Some(4)).is_err());
        assert!(checked_split_sizes(Some(1), 3, None, Some(3)).is_err());
        assert!(checked_split_sizes(Some(6), 2, Some(&[2, 3]), None).is_err());
    }

    #[test]
    fn flatten_axes_obey_schema_versions_and_scalar_bounds() {
        assert_eq!(checked_flatten_shape(&[], 0, false).unwrap(), [1, 1]);
        assert_eq!(
            checked_flatten_shape(&[2, 3, 4], -3, true).unwrap(),
            [1, 24]
        );
        assert_eq!(checked_flatten_shape(&[2, 3, 4], -1, true).unwrap(), [6, 4]);
        assert_eq!(checked_flatten_shape(&[2, 3, 4], 3, true).unwrap(), [24, 1]);
        assert!(checked_flatten_axis(3, -1, false).is_err());
        assert!(checked_flatten_axis(0, 1, true).is_err());
        for axis in [i64::MIN, i64::MAX, 1 << 32, -4, 4] {
            assert!(checked_flatten_axis(3, axis, true).is_err());
        }
    }

    #[test]
    fn flatten_products_are_zero_aware_and_never_narrow_overflow() {
        assert_eq!(
            checked_flatten_shape(&[65536, 65536, 0], 0, true).unwrap(),
            [1, 0]
        );
        assert_eq!(
            checked_flatten_shape(&[65536, 65536, 0], 1, true).unwrap(),
            [65536, 0]
        );
        assert_eq!(
            checked_flatten_shape(&[0, 65536, 65536], 2, true).unwrap(),
            [0, 65536]
        );
        assert_eq!(
            checked_flatten_shape(
                &[
                    0,
                    i64::from(i32::MAX),
                    i64::from(i32::MAX),
                    i64::from(i32::MAX)
                ],
                4,
                true
            )
            .unwrap(),
            [0, 1]
        );
        for shape in [vec![0, 65536, 65536], vec![0, 1 << 32], vec![0, -1]] {
            assert!(checked_flatten_shape(&shape, 1, true).is_err());
        }
    }

    #[test]
    fn integer_range_counts_preserve_full_width_bounds() {
        assert_eq!(
            integer_range_count(i64::MAX - 6, i64::MAX - 2, 1).unwrap(),
            4
        );
        assert_eq!(
            integer_range_count(i64::MIN, i64::MAX, i64::MAX).unwrap(),
            3
        );
        assert_eq!(
            integer_range_count(i64::MAX, i64::MIN, i64::MIN).unwrap(),
            2
        );
        assert_eq!(integer_range_count(4, 4, i64::MIN).unwrap(), 0);
        assert_eq!(integer_range_count(4, 8, -1).unwrap(), 0);
        assert_eq!(integer_range_count(8, 4, 1).unwrap(), 0);
        assert!(integer_range_count(0, 1, 0).is_err());
        assert!(integer_range_count(i64::MIN, i64::MAX, 1).is_err());
    }

    #[test]
    fn pad_geometry_preserves_original_coordinate_overlap() {
        for pads in [[-4, 2], [2, -4]] {
            let geometry = checked_pad_geometry(&[3], &[0], &pads, "constant").unwrap();
            assert_eq!(geometry.output, vec![1]);
            assert!(geometry.empty_source);
        }
        for pads in [[i64::MIN, i64::MAX], [i64::MAX, i64::MIN]] {
            let geometry = checked_pad_geometry(&[3], &[0], &pads, "constant").unwrap();
            assert_eq!(geometry.output, vec![2]);
            assert!(geometry.empty_source);
        }
        let geometry = checked_pad_geometry(&[3], &[0], &[-1, 2], "constant").unwrap();
        assert_eq!(geometry.start, vec![1]);
        assert_eq!(geometry.stop, vec![3]);
        assert_eq!(geometry.low, vec![0]);
        assert_eq!(geometry.high, vec![2]);
        assert_eq!(geometry.output, vec![4]);
    }

    #[test]
    fn pad_geometry_validates_axes_empty_and_native_bounds() {
        assert!(checked_pad_axes(2, &[-1, 1], 4).is_err());
        assert!(checked_pad_axes(2, &[i64::MIN], 2).is_err());
        assert!(checked_pad_axes(2, &[2], 2).is_err());
        assert!(checked_pad_axes(2, &[0], 3).is_err());
        assert!(checked_pad_geometry(&[3], &[0], &[-4, 0], "constant").is_err());
        assert!(checked_pad_geometry(&[3], &[0], &[i64::MAX, 0], "constant").is_err());
        assert!(checked_pad_geometry(&[3], &[0], &[-4, 2], "edge").is_err());
        assert!(checked_pad_geometry(&[0], &[0], &[1, 1], "wrap").is_err());
        assert!(checked_pad_geometry(&[1 << 30; 3], &[], &[], "constant").is_err());
        assert!(checked_pad_geometry(&[0, i64::from(i32::MAX) + 1], &[], &[], "constant").is_err());
        let empty = checked_pad_geometry(&[0, 3], &[1], &[1, 1], "constant").unwrap();
        assert_eq!(empty.output, vec![0, 5]);
    }

    #[test]
    fn rank_changes_preserve_axes_presence_and_scalar_geometry() {
        assert_eq!(
            checked_rank_change_axes(&[1, 2, 1], None, false, true).unwrap(),
            vec![0, 2]
        );
        assert!(
            checked_rank_change_axes(&[1, 2, 1], Some(&[]), false, true)
                .unwrap()
                .is_empty()
        );
        assert!(
            checked_rank_change_axes(&[], None, false, true)
                .unwrap()
                .is_empty()
        );
        assert!(checked_rank_change_axes(&[], None, true, true).is_err());
        assert_eq!(
            checked_rank_change_axes(&[], Some(&[-1]), true, true).unwrap(),
            vec![0]
        );
    }

    #[test]
    fn rank_changes_validate_full_width_axes_and_duplicates() {
        for axis in [i64::MIN, i64::MAX, 1i64 << 32, -4, 3] {
            assert!(checked_rank_change_axes(&[1, 2, 1], Some(&[axis]), false, true).is_err());
        }
        assert!(checked_rank_change_axes(&[1, 2, 1], Some(&[1]), false, true).is_err());
        assert!(checked_rank_change_axes(&[1, 2, 1], Some(&[-1]), false, false).is_err());
        assert_eq!(
            checked_rank_change_axes(&[1, 2, 1], Some(&[2, -3, 0]), false, true).unwrap(),
            vec![0, 2]
        );
        assert!(checked_rank_change_axes(&[2, 3], Some(&[0, -4]), true, true).is_err());
        assert_eq!(
            checked_rank_change_axes(&[2, 3], Some(&[-1, 0]), true, true).unwrap(),
            vec![0, 3]
        );
    }
}