use std::collections::{HashMap, HashSet};
use crate::engine::{MlxError, NodeDesc, TranslationContext, dim_i32, mlx_dtype_from_onnx};
use crate::mlx::{Array, VectorArray};
use crate::registry::{
ClaimPredicate, ClaimResult, CompileShapeSafety, K_ANY_OPSET, NodeView, OpHandler,
OpRegistration, OpRegistry, is_mlx_cpu_float, is_mlx_supported,
};
use crate::sys::mlx;
use crate::sys::ort;
use crate::{deny, require};
fn random_key(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<mlx::mlx_array, MlxError> {
match n.floats.get("seed") {
Some(&seed) => ctx.emit(|res, _| unsafe { mlx::mlx_random_key(res, seed as u64) }),
None => {
if ctx.shapeless_compile() || ctx.shape_keyed_compile() {
return Err(
"MLX: an unseeded random draw must consume the live default key in eager execution"
.to_string(),
);
}
Ok(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 = if dtype == mlx::mlx_dtype__MLX_BFLOAT16 {
let standard = ctx.emit(|res, s| unsafe {
mlx::mlx_random_normal(
res,
shape.as_ptr(),
shape.len(),
mlx::mlx_dtype__MLX_FLOAT32,
0.0,
1.0,
key,
s,
)
})?;
let scale = ctx.scalar_f32(scale);
let mean = ctx.scalar_f32(mean);
let cpu = crate::mlx::Stream::new_default_cpu();
let s = cpu.as_raw();
let affine = (|| {
let scaled =
ctx.emit(|res, _| unsafe { mlx::mlx_multiply(res, standard, scale, s) })?;
let shifted = ctx.emit(|res, _| unsafe { mlx::mlx_add(res, scaled, mean, s) })?;
let out = ctx.emit(|res, _| unsafe { mlx::mlx_astype(res, shifted, dtype, s) })?;
let mut outputs = VectorArray::new();
outputs.append(out);
crate::mlx::eval(&outputs)?;
Ok::<_, MlxError>(out)
})();
if unsafe { mlx::mlx_synchronize(s) } != 0 {
return Err(match affine {
Err(error) => {
format!("{error}; BF16 Normal CPU affine synchronization also failed")
}
Ok(_) => "BF16 Normal CPU affine synchronization failed".to_string(),
});
}
affine?
} else {
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_like_shape(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<Vec<i32>, MlxError> {
ctx.tensor_shape(&n.inputs[0])?
.into_iter()
.map(dim_i32)
.collect()
}
fn random_normal_like_op(ctx: &mut TranslationContext, n: &NodeDesc) -> Result<(), MlxError> {
let shape = random_like_shape(ctx, n)?;
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 shape = random_like_shape(ctx, n)?;
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(())
}
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, }
}
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 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)
|| (normal
&& node.since_version() >= 22
&& out.dtype
== ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16),
"output dtype must be float32 or float16 (also bfloat16 for Normal since schema 22), 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_input_supported(node: &NodeView, dtype: ort::ONNXTensorElementDataType) -> bool {
use ort::*;
if dtype == ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16
&& node.since_version() < 22
{
return false;
}
is_mlx_supported(dtype)
|| dtype == ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE
|| dtype == ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING
|| dtype == ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_COMPLEX64
|| dtype == ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_COMPLEX128
}
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!(
random_like_input_supported(node, inp.dtype),
"input dtype {} is not legal for this RandomLike schema",
crate::registry::ort_dtype_name(inp.dtype)
);
require!(
is_random_float(out.dtype)
|| (normal
&& node.since_version() >= 22
&& out.dtype
== ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16),
"output dtype must be float32 or float16 (also bfloat16 for Normal since schema 22), 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!(
node.input_rank_known(0) == Some(false) || shapes_compatible(&inp.shape, &out.shape),
"input/output shapes must be compatible, got {:?} -> {:?}",
inp.shape,
out.shape
);
require!(
inp.shape
.iter()
.chain(&out.shape)
.all(|&dim| dim <= i32::MAX as i64),
"RandomLike dimensions must fit MLX's native int32 geometry"
);
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)
|| inp.dtype == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE
|| (node.since_version() >= 22
&& inp.dtype
== ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16),
"probability input dtype must be float32, float16, float64, or bfloat16 since opset 22, got {}",
crate::registry::ort_dtype_name(inp.dtype)
);
require!(
is_mlx_supported(out.dtype)
|| out.dtype == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE,
"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();
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 && next == 1 {
return None;
}
Some(existing)
} else if existing == 1 {
if 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(())
}
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 random_compile_shape_safety(node: &NodeView) -> CompileShapeSafety {
if matches!(node.op_type().as_str(), "RandomNormal" | "RandomNormalLike")
&& node.output_info(0).is_some_and(|info| {
info.dtype == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16
})
{
CompileShapeSafety::EagerOnly {
reason: "BF16 Normal completes its F32 affine transform on the MLX CPU stream",
}
} else if node.float_attr_opt("seed").is_none() {
CompileShapeSafety::EagerOnly {
reason: "unseeded random draws consume the live default key",
}
} else if matches!(
node.op_type().as_str(),
"RandomNormalLike" | "RandomUniformLike"
) && node
.input_info(0)
.is_some_and(|info| !is_mlx_supported(info.dtype))
{
CompileShapeSafety::EagerOnly {
reason: "reads non-MLX RandomLike geometry without converting payloads",
}
} else {
CompileShapeSafety::ShapeKeyedOnly {
reason: crate::registry::MLX_RANDOM_BITS_SHAPE_REASON,
}
}
}
fn shape_keyed(
registry: &mut OpRegistry,
op_type: &'static str,
min_opset: i32,
handler: OpHandler,
claim: ClaimPredicate,
) {
registry.register_shape_classified(
OpRegistration {
domain: "",
op_type,
min_opset,
max_opset: K_ANY_OPSET,
handler,
claim,
},
random_compile_shape_safety,
);
}
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 bernoulli_double_tests {
use crate::mlx::{Array, Stream, VectorArray, eval};
use crate::sys::mlx;
fn native(call: impl FnOnce(*mut mlx::mlx_array) -> i32) -> Array {
let mut raw = unsafe { mlx::mlx_array_new() };
let status = call(&mut raw);
let output = Array::from_raw(raw);
assert_eq!(status, 0, "native Bernoulli test construction failed");
output
}
fn evaluate(array: &Array) {
let mut outputs = VectorArray::new();
outputs.append(array.as_raw());
eval(&outputs).expect("native Bernoulli test evaluation failed");
}
fn key(seed: u64) -> Array {
native(|res| unsafe { mlx::mlx_random_key(res, seed) })
}
fn split(key: &Array, stream: &Stream) -> (Array, Array) {
let mut state = unsafe { mlx::mlx_array_new() };
let mut draw = unsafe { mlx::mlx_array_new() };
let status =
unsafe { mlx::mlx_random_split(&mut state, &mut draw, key.as_raw(), stream.as_raw()) };
let result = (Array::from_raw(state), Array::from_raw(draw));
assert_eq!(status, 0, "native key split failed");
result
}
fn bits(key: &Array, shape: &[i32], stream: &Stream) -> Vec<u32> {
let output = native(|res| unsafe {
mlx::mlx_random_bits(
res,
shape.as_ptr(),
shape.len(),
4,
key.as_raw(),
stream.as_raw(),
)
});
evaluate(&output);
assert_eq!(output.dtype(), mlx::mlx_dtype__MLX_UINT32);
assert!(output.size() > 0);
let data = output.data_bytes().cast::<u32>();
assert!(!data.is_null());
unsafe { std::slice::from_raw_parts(data, output.size()).to_vec() }
}
fn sample(probabilities: &[f64], shape: &[i32], key: &Array, stream: &Stream) -> Vec<u8> {
let input = Array::from_data(
probabilities.as_ptr().cast(),
shape,
mlx::mlx_dtype__MLX_FLOAT64,
);
let output = native(|res| unsafe {
mlx::mlx_random_bernoulli(
res,
input.as_raw(),
shape.as_ptr(),
shape.len(),
key.as_raw(),
stream.as_raw(),
)
});
evaluate(&output);
assert_eq!(output.dtype(), mlx::mlx_dtype__MLX_BOOL);
assert_eq!(
output.shape(),
shape.iter().map(|&d| i64::from(d)).collect::<Vec<_>>()
);
assert_eq!(output.size(), probabilities.len());
if output.size() == 0 {
return Vec::new();
}
let data = output.data_bytes();
assert!(!data.is_null());
unsafe { std::slice::from_raw_parts(data, output.size()).to_vec() }
}
fn reference(bits: &[u32], probabilities: &[f64]) -> Vec<u8> {
assert_eq!(bits.len(), probabilities.len());
let upper = 4_294_967_296.0_f64 + 512.0;
bits.iter()
.zip(probabilities)
.map(|(&r, &p)| u8::from(f64::from(r) < p * upper))
.collect()
}
#[test]
fn bernoulli_preserves_double_thresholds_and_live_empty_default_keys() {
let stream = Stream::new_default_cpu();
let shape = [256];
for seed in [0, 41] {
let key = key(seed);
let draws = bits(&key, &shape, &stream);
let mut below = vec![0.5_f64; draws.len()];
let mut above = below.clone();
let mut straddles = Vec::new();
for (index, &r) in draws.iter().enumerate() {
let center = f64::from(r) / (4_294_967_296.0 + 512.0);
if !(0.25..0.75).contains(¢er) {
continue;
}
let low = f64::from_bits(center.to_bits() - 4);
let high = f64::from_bits(center.to_bits() + 4);
if (low as f32).to_bits() != (high as f32).to_bits() {
continue;
}
below[index] = low;
above[index] = high;
straddles.push(index);
}
assert!(
!straddles.is_empty(),
"fixture needs same-F32 F64 straddles"
);
let low = sample(&below, &shape, &key, &stream);
let high = sample(&above, &shape, &key, &stream);
assert_eq!(low, reference(&draws, &below));
assert_eq!(high, reference(&draws, &above));
for index in straddles {
assert_eq!((low[index], high[index]), (0, 1));
}
assert_eq!(low, sample(&below, &shape, &key, &stream));
for (p, expected) in [(0.0, 0), (1.0, 1)] {
assert_eq!(sample(&[p], &[], &key, &stream), vec![expected]);
assert_eq!(
sample(&vec![p; draws.len()], &shape, &key, &stream),
vec![expected; draws.len()]
);
}
assert!(sample(&[], &[0, 3], &key, &stream).is_empty());
}
let seed = 7919;
let (state, _) = split(&key(seed), &stream);
let (state, second_key) = split(&state, &stream);
let (_, third_key) = split(&state, &stream);
let probabilities = vec![0.5; 256];
let expected_second = reference(&bits(&second_key, &shape, &stream), &probabilities);
let expected_third = reference(&bits(&third_key, &shape, &stream), &probabilities);
assert_ne!(expected_second, expected_third);
assert_eq!(unsafe { mlx::mlx_random_seed(seed) }, 0);
let default_key = Array::new();
assert!(sample(&[], &[0, 3], &default_key, &stream).is_empty());
let explicit = key(41);
let fixed = sample(&probabilities, &shape, &explicit, &stream);
assert_eq!(fixed, sample(&probabilities, &shape, &explicit, &stream));
assert_eq!(
sample(&probabilities, &shape, &default_key, &stream),
expected_second
);
assert_eq!(
sample(&probabilities, &shape, &default_key, &stream),
expected_third
);
assert_eq!(unsafe { mlx::mlx_synchronize(stream.as_raw()) }, 0);
}
}
#[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));
}
}