onnxruntime-ep-mlx 0.29.6

//! Random / miscellaneous op handlers (ai.onnx): RandomNormal(+Like), RandomUniform(+Like),
//! Bernoulli, Multinomial, Einsum. Faithful port of the C++ `ops/randommisc.cc`. NonZero / Unique are
//! deliberately left to ORT CPU (mlx-c has no nonzero/unique primitive).

use std::collections::{HashMap, HashSet};

use crate::engine::{MlxError, NodeDesc, TranslationContext, mlx_dtype_from_onnx};
use crate::mlx::{Array, VectorArray};
use crate::registry::{
    ClaimPredicate, ClaimResult, K_ANY_OPSET, NodeView, OpHandler, OpRegistration, OpRegistry,
    is_mlx_cpu_float, is_mlx_float, is_mlx_supported,
};
use crate::sys::mlx;
use crate::sys::ort;
use crate::{deny, require};

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

/// The MLX PRNG key from a `seed` float attribute, or an empty array (default key) when absent.
fn random_key(ctx: &mut TranslationContext, n: &NodeDesc) -> mlx::mlx_array {
    match n.floats.get("seed") {
        Some(&seed) => {
            let raw = unsafe {
                let mut r = mlx::mlx_array_new();
                mlx::mlx_random_key(&mut r, seed as u64);
                r
            };
            ctx.keep(Array::from_raw(raw))
        }
        None => ctx.keep(Array::new()),
    }
}

fn attr_shape(n: &NodeDesc) -> Vec<i32> {
    n.int_arrays
        .get("shape")
        .map(|v| v.iter().map(|&d| d as i32).collect())
        .unwrap_or_default()
}

fn random_normal_with_shape(
    ctx: &mut TranslationContext,
    n: &NodeDesc,
    shape: Vec<i32>,
) -> Result<(), MlxError> {
    let key = random_key(ctx, n);
    let dtype = mlx_dtype_from_onnx(n.outputs[0].otype);
    let mean = n.floats.get("mean").copied().unwrap_or(0.0);
    let scale = n.floats.get("scale").copied().unwrap_or(1.0);
    let out = ctx.emit(|res, s| unsafe {
        mlx::mlx_random_normal(res, shape.as_ptr(), shape.len(), dtype, mean, scale, key, s)
    })?;
    ctx.bind(&n.outputs[0], out);
    Ok(())
}

fn random_normal_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    random_normal_with_shape(ctx, n, attr_shape(n))
}

fn random_normal_like_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let x = ctx.resolve(&n.inputs[0])?;
    let shape = ctx.shape_of(x);
    random_normal_with_shape(ctx, n, shape)
}

fn random_uniform_with_shape(
    ctx: &mut TranslationContext,
    n: &NodeDesc,
    shape: Vec<i32>,
) -> Result<(), MlxError> {
    let low = ctx.scalar_f32(n.floats.get("low").copied().unwrap_or(0.0));
    let high = ctx.scalar_f32(n.floats.get("high").copied().unwrap_or(1.0));
    let key = random_key(ctx, n);
    let dtype = mlx_dtype_from_onnx(n.outputs[0].otype);
    let out = ctx.emit(|res, s| unsafe {
        mlx::mlx_random_uniform(res, low, high, shape.as_ptr(), shape.len(), dtype, key, s)
    })?;
    ctx.bind(&n.outputs[0], out);
    Ok(())
}

fn random_uniform_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    random_uniform_with_shape(ctx, n, attr_shape(n))
}

fn random_uniform_like_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let x = ctx.resolve(&n.inputs[0])?;
    let shape = ctx.shape_of(x);
    random_uniform_with_shape(ctx, n, shape)
}

fn bernoulli_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let probs = ctx.resolve(&n.inputs[0])?;
    let shape = ctx.shape_of(probs);
    let key = random_key(ctx, n);
    let sampled = ctx.emit(|res, s| unsafe {
        mlx::mlx_random_bernoulli(res, probs, shape.as_ptr(), shape.len(), key, s)
    })?;
    let out = ctx.astype(sampled, mlx_dtype_from_onnx(n.outputs[0].otype))?;
    ctx.bind(&n.outputs[0], out);
    Ok(())
}

fn multinomial_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let logits = ctx.resolve(&n.inputs[0])?;
    let sample_size = n.ints.get("sample_size").copied().unwrap_or(1) as i32;
    let key = random_key(ctx, n);
    let sampled = ctx.emit(|res, s| unsafe {
        mlx::mlx_random_categorical_num_samples(res, logits, -1, sample_size, key, s)
    })?;
    let out = ctx.astype(sampled, mlx_dtype_from_onnx(n.outputs[0].otype))?;
    ctx.bind(&n.outputs[0], out);
    Ok(())
}

fn einsum_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
    let mut operands = VectorArray::new();
    for input in &n.inputs {
        let a = ctx.resolve(input)?;
        operands.append(a);
    }
    let equation: String = n
        .strings
        .get("equation")
        .cloned()
        .unwrap_or_default()
        .chars()
        .filter(|c| !c.is_whitespace())
        .collect();
    let ceq = std::ffi::CString::new(equation).map_err(|_| "einsum: bad equation".to_string())?;
    let operands_raw = operands.as_raw();
    let out = ctx.emit(|res, s| unsafe { mlx::mlx_einsum(res, ceq.as_ptr(), operands_raw, s) })?;
    ctx.bind(&n.outputs[0], out);
    Ok(())
}

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

/// Optional-seed validation: absent is fine; a present seed must be a finite non-negative float.
fn optional_seed_supported(node: &NodeView) -> bool {
    if !node.has_attr("seed") {
        return true;
    }
    match node.float_attr_opt("seed") {
        Some(seed) => seed.is_finite() && seed >= 0.0 && (seed as f64) < 2f64.powi(64),
        None => false, // present but not a float
    }
}

fn is_random_float(t: ort::ONNXTensorElementDataType) -> bool {
    t == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT
        || t == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16
}

fn is_boundary_type(t: ort::ONNXTensorElementDataType) -> bool {
    use ort::*;
    is_mlx_float(t)
        || t == ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL
        || t == ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8
        || t == ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16
        || t == ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32
        || t == ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64
        || t == ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8
        || t == ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT16
        || t == ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT32
}

fn valid_shape(shape: &[i64]) -> bool {
    shape.iter().all(|&d| d >= 0 && d <= i32::MAX as i64)
}

fn shapes_compatible(a: &[i64], b: &[i64]) -> bool {
    if a.len() != b.len() {
        return false;
    }
    a.iter()
        .zip(b.iter())
        .all(|(&x, &y)| x < 0 || y < 0 || x == y)
}

fn random_shape_claim(node: &NodeView, normal: bool) -> ClaimResult {
    require!(
        node.num_inputs() == 0 && node.num_outputs() == 1,
        "expects 0 inputs and 1 output, got {}in/{}out",
        node.num_inputs(),
        node.num_outputs()
    );
    require!(
        optional_seed_supported(node),
        "seed must be a finite non-negative float representable as u64"
    );
    let out = match node.output_info(0) {
        Some(o) => o,
        None => deny!("missing tensor type/shape info on output"),
    };
    require!(
        is_random_float(out.dtype),
        "output dtype must be float32 or float16, got {}",
        crate::registry::ort_dtype_name(out.dtype)
    );
    let dtype_attr = node.int_attr(
        "dtype",
        ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT as i64,
    );
    require!(
        dtype_attr == out.dtype as i64,
        "dtype attribute {} must match output dtype {}",
        crate::registry::ort_dtype_name(dtype_attr as ort::ONNXTensorElementDataType),
        crate::registry::ort_dtype_name(out.dtype)
    );
    let (present, attr_shape) = node.ints_attr("shape");
    require!(present, "shape attribute is required");
    require!(
        valid_shape(&attr_shape),
        "shape dimensions must be in 0..=i32::MAX (got {:?})",
        attr_shape
    );
    require!(
        out.shape == attr_shape,
        "shape attribute {:?} must match output shape {:?}",
        attr_shape,
        out.shape
    );
    if normal {
        let mean = node.float_attr_opt("mean").unwrap_or(0.0);
        let scale = node.float_attr_opt("scale").unwrap_or(1.0);
        require!(
            mean.is_finite() && scale.is_finite() && scale >= 0.0,
            "mean must be finite and scale finite/non-negative (got mean={mean}, scale={scale})"
        );
    } else {
        let low = node.float_attr_opt("low").unwrap_or(0.0);
        let high = node.float_attr_opt("high").unwrap_or(1.0);
        require!(
            low.is_finite() && high.is_finite() && low < high,
            "low/high must be finite with low < high (got low={low}, high={high})"
        );
    }
    Ok(())
}

fn random_normal_claim(node: &NodeView) -> ClaimResult {
    random_shape_claim(node, true)
}

fn random_uniform_claim(node: &NodeView) -> ClaimResult {
    random_shape_claim(node, false)
}

fn random_like_claim(node: &NodeView, normal: bool) -> 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()
    );
    require!(
        optional_seed_supported(node),
        "seed must be a finite non-negative float representable as u64"
    );
    let (inp, out) = match (node.input_info(0), node.output_info(0)) {
        (Some(a), Some(b)) => (a, b),
        _ => deny!("missing tensor type/shape info on input or output"),
    };
    require!(
        is_mlx_supported(inp.dtype),
        "input dtype {} is not supported by MLX",
        crate::registry::ort_dtype_name(inp.dtype)
    );
    require!(
        is_random_float(out.dtype),
        "output dtype must be float32 or float16, got {}",
        crate::registry::ort_dtype_name(out.dtype)
    );
    let dtype_attr = node.int_attr("dtype", inp.dtype as i64);
    require!(
        dtype_attr == out.dtype as i64,
        "dtype attribute {} must match output dtype {}",
        crate::registry::ort_dtype_name(dtype_attr as ort::ONNXTensorElementDataType),
        crate::registry::ort_dtype_name(out.dtype)
    );
    require!(
        shapes_compatible(&inp.shape, &out.shape),
        "input/output shapes must be compatible, got {:?} -> {:?}",
        inp.shape,
        out.shape
    );
    if normal {
        let mean = node.float_attr_opt("mean").unwrap_or(0.0);
        let scale = node.float_attr_opt("scale").unwrap_or(1.0);
        require!(
            mean.is_finite() && scale.is_finite() && scale >= 0.0,
            "mean must be finite and scale finite/non-negative (got mean={mean}, scale={scale})"
        );
    } else {
        let low = node.float_attr_opt("low").unwrap_or(0.0);
        let high = node.float_attr_opt("high").unwrap_or(1.0);
        require!(
            low.is_finite() && high.is_finite() && low < high,
            "low/high must be finite with low < high (got low={low}, high={high})"
        );
    }
    Ok(())
}

fn random_normal_like_claim(node: &NodeView) -> ClaimResult {
    random_like_claim(node, true)
}

fn random_uniform_like_claim(node: &NodeView) -> ClaimResult {
    random_like_claim(node, false)
}

fn bernoulli_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()
    );
    require!(
        optional_seed_supported(node),
        "seed must be a finite non-negative float representable as u64"
    );
    let (inp, out) = match (node.input_info(0), node.output_info(0)) {
        (Some(a), Some(b)) => (a, b),
        _ => deny!("missing tensor type/shape info on input or output"),
    };
    require!(
        is_random_float(inp.dtype),
        "probability input dtype must be float32 or float16, got {}",
        crate::registry::ort_dtype_name(inp.dtype)
    );
    require!(
        is_boundary_type(out.dtype),
        "output dtype {} is unsupported",
        crate::registry::ort_dtype_name(out.dtype)
    );
    let dtype_attr = node.int_attr("dtype", inp.dtype as i64);
    require!(
        dtype_attr == out.dtype as i64,
        "dtype attribute {} must match output dtype {}",
        crate::registry::ort_dtype_name(dtype_attr as ort::ONNXTensorElementDataType),
        crate::registry::ort_dtype_name(out.dtype)
    );
    require!(
        shapes_compatible(&inp.shape, &out.shape),
        "input/output shapes must be compatible, got {:?} -> {:?}",
        inp.shape,
        out.shape
    );
    Ok(())
}

fn multinomial_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()
    );
    require!(
        optional_seed_supported(node),
        "seed must be a finite non-negative float representable as u64"
    );
    let (inp, out) = match (node.input_info(0), node.output_info(0)) {
        (Some(a), Some(b)) => (a, b),
        _ => deny!("missing tensor type/shape info on input or output"),
    };
    let sample_size = node.int_attr("sample_size", 1);
    require!(
        is_random_float(inp.dtype),
        "input dtype must be float32 or float16, got {}",
        crate::registry::ort_dtype_name(inp.dtype)
    );
    require!(
        out.dtype == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32
            || out.dtype == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64,
        "output dtype must be int32 or int64, got {}",
        crate::registry::ort_dtype_name(out.dtype)
    );
    let dtype_attr = node.int_attr(
        "dtype",
        ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32 as i64,
    );
    require!(
        dtype_attr == out.dtype as i64,
        "dtype attribute {} must match output dtype {}",
        crate::registry::ort_dtype_name(dtype_attr as ort::ONNXTensorElementDataType),
        crate::registry::ort_dtype_name(out.dtype)
    );
    require!(
        inp.shape.len() == 2 && out.shape.len() == 2,
        "input/output must both have rank 2, got rank {} -> {}",
        inp.shape.len(),
        out.shape.len()
    );
    require!(
        inp.shape[1] > 0,
        "class dimension must be static and positive (got {})",
        inp.shape[1]
    );
    require!(
        sample_size > 0 && sample_size <= i32::MAX as i64,
        "sample_size must be in 1..=i32::MAX (got {sample_size})"
    );
    require!(
        inp.shape[0] < 0 || out.shape[0] < 0 || inp.shape[0] == out.shape[0],
        "batch dimensions must match, got {} -> {}",
        inp.shape[0],
        out.shape[0]
    );
    require!(
        out.shape[1] < 0 || out.shape[1] == sample_size,
        "output sample dimension must equal sample_size {sample_size}, got {}",
        out.shape[1]
    );
    Ok(())
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum EinsumToken {
    Label(char),
    Ellipsis,
}

fn parse_einsum_term(raw: &str, output: bool) -> Option<Vec<EinsumToken>> {
    let bytes = raw.as_bytes();
    let mut tokens = Vec::new();
    let mut labels = HashSet::new();
    let mut has_ellipsis = false;
    let mut i = 0;
    while i < bytes.len() {
        let c = bytes[i] as char;
        if c.is_ascii_alphabetic() {
            if output && !labels.insert(c) {
                return None;
            }
            tokens.push(EinsumToken::Label(c));
            i += 1;
        } else if bytes[i..].starts_with(b"...") && !has_ellipsis {
            tokens.push(EinsumToken::Ellipsis);
            has_ellipsis = true;
            i += 3;
        } else {
            return None;
        }
    }
    Some(tokens)
}

fn parse_einsum(raw: &str) -> Option<(Vec<Vec<EinsumToken>>, Vec<EinsumToken>)> {
    let eq: String = raw.chars().filter(|c| !c.is_whitespace()).collect();
    // MLX 0.32.2 includes equation punctuation in its internal used-character set before reserving
    // virtual labels. More than 52 distinct raw characters underflows that reservation.
    if eq.chars().collect::<HashSet<_>>().len() > 52 {
        return None;
    }
    let (lhs, explicit_output) = match eq.split_once("->") {
        Some((lhs, output)) if !output.contains("->") => (lhs, Some(output)),
        Some(_) => return None,
        None => (eq.as_str(), None),
    };
    if lhs.is_empty() {
        return None;
    }
    let terms = lhs
        .split(',')
        .map(|term| parse_einsum_term(term, false))
        .collect::<Option<Vec<_>>>()?;
    let output = if let Some(output) = explicit_output {
        parse_einsum_term(output, true)?
    } else {
        let mut counts = HashMap::new();
        let mut has_ellipsis = false;
        for token in terms.iter().flatten() {
            match token {
                EinsumToken::Label(label) => *counts.entry(*label).or_insert(0usize) += 1,
                EinsumToken::Ellipsis => has_ellipsis = true,
            }
        }
        let mut labels = counts
            .into_iter()
            .filter_map(|(label, count)| (count == 1).then_some(label))
            .collect::<Vec<_>>();
        labels.sort_unstable();
        let mut output = Vec::with_capacity(labels.len() + usize::from(has_ellipsis));
        if has_ellipsis {
            output.push(EinsumToken::Ellipsis);
        }
        output.extend(labels.into_iter().map(EinsumToken::Label));
        output
    };
    Some((terms, output))
}

fn merge_broadcast_dim(existing: i64, next: i64) -> Option<i64> {
    if existing < 0 && next < 0 {
        None
    } else if existing == next || next == 1 {
        if (existing == 0 || existing < 0) && next == 1 {
            return None;
        }
        Some(existing)
    } else if existing == 1 {
        if next == 0 || next < 0 {
            return None;
        }
        Some(next)
    } else if existing == 0 || next == 0 {
        None
    } else if existing < 0 {
        Some(next)
    } else if next < 0 {
        Some(existing)
    } else {
        None
    }
}

fn einsum_claim(node: &NodeView) -> ClaimResult {
    let ni = node.num_inputs();
    require!(
        ni > 0 && node.num_outputs() == 1,
        "expects at least 1 input and exactly 1 output, got {}in/{}out",
        ni,
        node.num_outputs()
    );
    let equation = node.string_attr("equation", "");
    require!(
        node.has_attr("equation") && !equation.is_empty(),
        "non-empty equation attribute is required"
    );
    let (input_terms, output_term) = match parse_einsum(&equation) {
        Some(v) => v,
        None => deny!(
            "equation contains invalid labels, ellipsis, arrow, or repeated output labels (got {equation:?})"
        ),
    };
    require!(
        input_terms.len() == ni,
        "equation has {} input terms but node has {ni} inputs",
        input_terms.len()
    );
    let (in0, out) = match (node.input_info(0), node.output_info(0)) {
        (Some(a), Some(b)) => (a, b),
        _ => deny!("missing tensor type/shape info on first input or output"),
    };
    let dtype = in0.dtype;
    require!(
        is_mlx_cpu_float(dtype) && out.dtype == dtype,
        "all tensors must share an MLX fp32/fp16/bf16/fp64 dtype; v28 integer, float8, and sub-byte Einsum forms stay on CPU because MLX 0.32.2's native Einsum path does not execute integer contractions safely, got first input {} and output {}",
        crate::registry::ort_dtype_name(dtype),
        crate::registry::ort_dtype_name(out.dtype)
    );
    let mut dims: HashMap<char, i64> = HashMap::new();
    let mut ellipsis_dims: Vec<i64> = Vec::new();
    let mut input_has_ellipsis = false;
    let mut input_ellipsis_rank: Option<usize> = None;
    for (i, term) in input_terms.iter().enumerate() {
        let info = match node.input_info(i) {
            Some(x) => x,
            None => deny!("missing tensor type/shape info on input {i}"),
        };
        require!(
            info.dtype == dtype,
            "input {i} dtype must match {}, got {}",
            crate::registry::ort_dtype_name(dtype),
            crate::registry::ort_dtype_name(info.dtype)
        );
        let label_count = term
            .iter()
            .filter(|token| matches!(token, EinsumToken::Label(_)))
            .count();
        let has_ellipsis = term.contains(&EinsumToken::Ellipsis);
        input_has_ellipsis |= has_ellipsis;
        require!(
            (has_ellipsis && info.shape.len() >= label_count)
                || (!has_ellipsis && info.shape.len() == label_count),
            "input {i} rank {} is incompatible with {} labels{}",
            info.shape.len(),
            label_count,
            if has_ellipsis { " plus ellipsis" } else { "" }
        );
        let ellipsis_rank = info.shape.len() - label_count;
        if has_ellipsis {
            match input_ellipsis_rank {
                Some(expected) => require!(
                    expected == ellipsis_rank,
                    "input {i} ellipsis rank {ellipsis_rank} must match prior ellipsis rank {expected}"
                ),
                None => {
                    input_ellipsis_rank = Some(ellipsis_rank);
                    ellipsis_dims = vec![1; ellipsis_rank];
                }
            }
        }
        let mut axis = 0;
        let mut local_dims = HashMap::new();
        for token in term {
            let EinsumToken::Label(label) = token else {
                let start = ellipsis_dims.len() - ellipsis_rank;
                for (offset, &d) in info.shape[axis..axis + ellipsis_rank].iter().enumerate() {
                    let existing = ellipsis_dims[start + offset];
                    let Some(merged) = merge_broadcast_dim(existing, d) else {
                        deny!("input {i} ellipsis has incompatible dimensions {existing} and {d}")
                    };
                    ellipsis_dims[start + offset] = merged;
                }
                axis += ellipsis_rank;
                continue;
            };
            let d = info.shape[axis];
            if let Some(existing) = local_dims.insert(*label, d) {
                require!(
                    existing < 0 || d < 0 || existing == d,
                    "input {i} repeated label {label:?} has incompatible dimensions {existing} and {d}"
                );
            }
            match dims.get(label).copied() {
                None => {
                    dims.insert(*label, d);
                }
                Some(existing) => {
                    let Some(merged) = merge_broadcast_dim(existing, d) else {
                        deny!("label {label:?} has incompatible dimensions {existing} and {d}")
                    };
                    dims.insert(*label, merged);
                }
            }
            axis += 1;
        }
    }
    let unique_labels = dims.len();
    let ellipsis_rank = input_ellipsis_rank.unwrap_or(0);
    require!(
        unique_labels + ellipsis_rank <= 52,
        "equation needs {} labels after ellipsis expansion, but MLX supports at most 52",
        unique_labels + ellipsis_rank
    );
    let output_has_ellipsis = output_term.contains(&EinsumToken::Ellipsis);
    require!(
        ellipsis_rank == 0 || output_has_ellipsis,
        "a non-empty input ellipsis must appear in the output"
    );
    let mut expected_output = Vec::new();
    for token in output_term {
        match token {
            EinsumToken::Ellipsis => {
                require!(
                    input_has_ellipsis,
                    "output ellipsis requires an ellipsis in at least one input term"
                );
                expected_output.extend_from_slice(&ellipsis_dims);
            }
            EinsumToken::Label(label) => match dims.get(&label).copied() {
                Some(d) => expected_output.push(d),
                None => deny!("output label {label:?} does not appear in any input term"),
            },
        }
    }
    require!(
        out.shape.len() == expected_output.len(),
        "output rank {} must match inferred equation rank {}",
        out.shape.len(),
        expected_output.len()
    );
    for (axis, (&expected, &actual)) in expected_output.iter().zip(&out.shape).enumerate() {
        require!(
            expected < 0 || actual < 0 || expected == actual,
            "output axis {axis} dimension {actual} does not match inferred dimension {expected}"
        );
    }
    Ok(())
}

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

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

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

pub fn register(registry: &mut OpRegistry) {
    shape_keyed(
        registry,
        "RandomNormal",
        1,
        random_normal_op,
        random_normal_claim,
    );
    shape_keyed(
        registry,
        "RandomNormalLike",
        1,
        random_normal_like_op,
        random_normal_like_claim,
    );
    shape_keyed(
        registry,
        "RandomUniform",
        1,
        random_uniform_op,
        random_uniform_claim,
    );
    shape_keyed(
        registry,
        "RandomUniformLike",
        1,
        random_uniform_like_op,
        random_uniform_like_claim,
    );
    shape_keyed(registry, "Bernoulli", 15, bernoulli_op, bernoulli_claim);
    shape_keyed(
        registry,
        "Multinomial",
        7,
        multinomial_op,
        multinomial_claim,
    );
    shapeless(registry, "Einsum", 12, einsum_op, einsum_claim);
}

#[cfg(test)]
mod einsum_parser_tests {
    use super::{merge_broadcast_dim, parse_einsum};

    #[test]
    fn scalar_terms_and_outputs_are_valid() {
        assert!(parse_einsum(",ij->ij").is_some());
        assert!(parse_einsum("ij,,jk->ik").is_some());
        assert!(parse_einsum("i,i->").is_some());
        assert!(parse_einsum("ij,->ij").is_some());
    }

    #[test]
    fn mlx_raw_character_limit_is_enforced() {
        let labels = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ";
        assert!(parse_einsum(&format!("{labels}->{labels}")).is_none());
        let safe = &labels[..50];
        assert!(parse_einsum(&format!("{safe}->{safe}")).is_some());
    }

    #[test]
    fn mlx_unsafe_zero_size_broadcast_possibilities_are_declined() {
        assert_eq!(merge_broadcast_dim(0, 1), None);
        assert_eq!(merge_broadcast_dim(1, 0), None);
        assert_eq!(merge_broadcast_dim(-1, -1), None);
        assert_eq!(merge_broadcast_dim(-1, 1), None);
        assert_eq!(merge_broadcast_dim(1, -1), None);
        assert_eq!(merge_broadcast_dim(-1, 4), Some(4));
        assert_eq!(merge_broadcast_dim(4, -1), Some(4));
    }
}