Enum Primitive 

#[non_exhaustive]
pub enum Primitive {
    Unary(UnaryOp),
    Binary(BinaryOp),
    Select,
    AsType(DType),
    Reshape(Vec<usize>),
    Squeeze(Vec<usize>),
    ExpandDims(Vec<usize>),
    Transpose(Vec<usize>),
    Slice {
        start: Vec<usize>,
        step: Vec<isize>,
        len: Vec<usize>,
    },
    Broadcast(Vec<usize>),
    StopGradient,
    View(DType),
    ToFp8,
    FromFp8(DType),
    SortedOffsets {
        groups: usize,
    },
    SliceUpdate {
        start: Vec<usize>,
        step: Vec<isize>,
        len: Vec<usize>,
    },
    Reduce {
        op: ReduceOp,
        axes: Vec<usize>,
    },
    ArgReduce {
        op: ArgReduceOp,
        axis: usize,
    },
    GatherAxis {
        axis: usize,
    },
    ScatterAxis {
        axis: usize,
        reduce: ScatterReduce,
    },
    Sort {
        axis: usize,
    },
    ArgSort {
        axis: usize,
    },
    Scan {
        op: ScanOp,
        axis: usize,
        reverse: bool,
        inclusive: bool,
    },
    Concatenate {
        axis: usize,
    },
    Matmul,
    AddMM {
        alpha: u32,
        beta: u32,
    },
    Quantize {
        quant: Quant,
    },
    Dequantize {
        quant: Quant,
        dtype: DType,
    },
    QuantizedMatmul {
        quant: Quant,
        transpose: bool,
    },
    Fast(FastOp),
    Compiled(Fused),
    RandomBits {
        shape: Vec<usize>,
        width: u8,
    },
    Custom(CustomCall),
    MetalKernel(KernelCall),
    Copy,
    Depends {
        n_deps: usize,
    },
    Collective {
        op: CollectiveOp,
        group: usize,
    },
}

An operation recorded in the graph, with its parameters.

Variants (Non-exhaustive)

Non-exhaustive enums could have additional variants added in future. Therefore, when matching against variants of non-exhaustive enums, an extra wildcard arm must be added to account for any future variants.

Unary(UnaryOp)

Element-wise function of one operand.

Binary(BinaryOp)

Element-wise function of two operands of equal shape and dtype.

Select

where(condition, a, b): a boolean condition selects from a or b.

AsType(DType)

Conversion to another element type.

Reshape(Vec<usize>)

The same elements in another shape. A view when the layout permits, else a copy.

Squeeze(Vec<usize>)

The input without its size-1 dimensions at axes (sorted, unique). A view. Unlike Reshape, the output shape follows the input’s (MLX’s Squeeze, which shapeless compile replays on any shape).

ExpandDims(Vec<usize>)

The input with size-1 dimensions inserted at the output’s axes (sorted, unique). A view (MLX’s ExpandDims).

Transpose(Vec<usize>)

Axes reordered: output axis i is input axis axes[i]. A view.

Slice

A normalised slice per dimension. A view.

Fields

start: Vec<usize>

First index per dimension.

step: Vec<isize>

Step per dimension (non-zero, may be negative).

len: Vec<usize>

Number of selected indices per dimension (the output shape).

Broadcast(Vec<usize>)

Presented as a larger shape with zero strides. A view.

StopGradient

The input unchanged, but no gradient flows through it. A view.

View(DType)

The input’s bytes reinterpreted as dtype (MLX’s view): the same shape when the element sizes agree, else the last dimension scaled by their ratio. A view of a row-contiguous input (or of equal element sizes); otherwise a copy, reinterpreted.

ToFp8

Floats as E4M3 bytes (uint8), rounding to nearest even and saturating at ±448 (MLX’s to_fp8).

FromFp8(DType)

E4M3 bytes (uint8) as dtype floats (MLX’s from_fp8).

SortedOffsets

Sorted uint32 indices [..., M] to int32 [..., groups]: where each group’s rows start along the last axis (the lower bound of g; MLX’s gather_mm_offsets), which the grouped gather kernels walk.

Fields

groups: usize

The number of groups (weight batches).

SliceUpdate

(base, update): a copy of base whose region selected as by Slice is replaced by update (of shape len).

Fields

start: Vec<usize>

First index per dimension.

step: Vec<isize>

Step per dimension (non-zero, may be negative).

len: Vec<usize>

Number of selected indices per dimension (the update’s shape).

Reduce

Reduction over sorted, unique axes; reduced dimensions become size 1.

Fields

op: ReduceOp

The reduction.

axes: Vec<usize>

Reduced axes, sorted and unique.

ArgReduce

Index of an extreme value along axis (uint32); the axis becomes size 1.

Fields

op: ArgReduceOp

Maximum or minimum.

axis: usize

The reduced axis.

GatherAxis

(src, indices): the elements of src at indices along axis (MLX’s GatherAxis, take_along_axis). indices (integers) has src’s shape except along axis and gives the output’s shape. Negative indices count from the end; out-of-range indices are clamped.

Fields

axis: usize

The indexed axis.

ScatterAxis

(dst, indices, updates): a copy of dst with updates placed at indices along axis (MLX’s ScatterAxis, put_along_axis). indices (integers) and updates share a shape, which is dst’s except along axis. Indices as in GatherAxis.

Fields

axis: usize

The indexed axis.

reduce: ScatterReduce

How updates combine with dst.

Sort

The input sorted along axis (ascending, stable, NaN last).

Fields

axis: usize

The sorted axis.

ArgSort

The indices (uint32) that sort the input along axis, as Primitive::Sort.

Fields

axis: usize

The sorted axis.

Scan

A cumulative scan along axis.

Fields

op: ScanOp

The combination.

axis: usize

The scanned axis.

reverse: bool

Scan from the end of the axis.

inclusive: bool

Whether each output includes its own element (else it holds the elements before it, starting from the combination’s identity).

Concatenate

The inputs (one dtype, equal shapes except along axis) joined along axis.

Fields

axis: usize

The joined axis.

Matmul

Batched matrix product of [B..., M, K] and [B..., K, N] with equal batch dimensions.

AddMM

alpha · (a @ b) + beta · c for inputs (a, b, c): a Matmul whose c has the output’s shape. The factors are stored as f32 bits (so the enum stays Eq + Hash).

Fields

alpha: u32

f32::to_bits of alpha.

beta: u32

f32::to_bits of beta.

Quantize

Quantization of the last axis in groups (MLX’s quantize). Outputs the packed codes (uint32, bits per value) and the scales: for affine, scales and biases in the input’s dtype; for the floating-point modes, uint8 scale codes.

Fields

quant: Quant

The mode, group size and bits.

Dequantize

(w, scales[, biases]): the packed values restored in dtype (affine: code · scale + bias in the scales’ dtype, which dtype equals; the floating-point modes have no biases).

Fields

quant: Quant

The mode, group size and bits.

dtype: DType

The output dtype.

QuantizedMatmul

(x, w, scales[, biases]): x · dequantize(w)ᵀ (transpose) or x · dequantize(w). w is two-dimensional, or has x’s batch dimensions. Biases exist for affine only.

Fields

quant: Quant

The mode, group size and bits.

transpose: bool

Whether the dequantized matrix is used transposed (MLX’s default).

Fast(FastOp)

A fused operation with a dedicated GPU kernel (crate::fast).

Compiled(Fused)

Element-wise operations fused into one generated kernel (crate::graph::Fused), as compile builds them. The inputs share the output’s shape.

RandomBits

Random bits from threefry2x32 keys of shape [..., 2] (uint32): the output has shape keys.shape[..-1] + shape and width bytes per element (MLX’s layout).

Fields

shape: Vec<usize>

The shape drawn per key.

width: u8

Bytes per element: 1, 2 or 4.

Custom(CustomCall)

A call of a crate::custom_function: the inputs are the primals followed by the function’s outputs, which this node presents unchanged (views); transforms use the function’s custom rules.

MetalKernel(KernelCall)

A custom Metal kernel (crate::fast::MetalKernel).

Copy

A row-contiguous copy of the input into a new buffer. Never simplified away (as a same-type conversion or identity view would be).

Depends

(x, deps...): x unchanged (sharing its data), evaluated only after the n_deps other inputs. Gradients flow to x only. It orders work, as MLX’s depends: a recomputation that depends on cotangents cannot be merged with the forward pass.

Fields

n_deps: usize

The number of dependencies after x.

Collective

A collective over a distributed group (by id; crate::distributed).

Fields

op: CollectiveOp

The operation and its parameters.

group: usize

The group’s id.

Implementations

impl Primitive

pub fn name(&self) -> &'static str

The operation’s name, as used in error messages.

pub fn fixes_shape(&self) -> bool

Whether the parameters hold sizes taken from the traced shapes, so a replay on other shapes would be wrong (MLX’s primitives without output_shapes): shapeless compile rejects these. Broadcast is listed; shapeless compile re-derives its shape from the operands it was broadcast against. Attention keeps kernel choices made for the traced lengths.

Trait Implementations

impl Clone for Primitive

fn clone(&self) -> Self

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for Primitive

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl Eq for Primitive

impl Hash for Primitive

fn hash<__H: Hasher>(&self, state: &mut __H)

Feeds this value into the given Hasher.

fn hash_slice<H>(data: &[Self], state: &mut H)

where H: Hasher, Self: Sized,

Feeds a slice of this type into the given Hasher.

impl PartialEq for Primitive

fn eq(&self, other: &Self) -> bool

Equality operator ==.

fn ne(&self, other: &Rhs) -> bool

Inequality operator !=.

impl StructuralPartialEq for Primitive