Enum Primitive 

#[non_exhaustive]
pub enum Primitive {
    Unary(UnaryOp),
    Binary(BinaryOp),
    Select,
    AsType(DType),
    Reshape(Vec<usize>),
    Transpose(Vec<usize>),
    Slice {
        start: Vec<usize>,
        step: Vec<isize>,
        len: Vec<usize>,
    },
    Broadcast(Vec<usize>),
    StopGradient,
    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 {
        group_size: usize,
        bits: u8,
    },
    Dequantize {
        group_size: usize,
        bits: u8,
    },
    QuantizedMatmul {
        group_size: usize,
        bits: u8,
        transpose: bool,
    },
    Fast(FastOp),
    Compiled(Fused),
    RandomBits {
        shape: Vec<usize>,
        width: u8,
    },
    Custom(CustomCall),
    MetalKernel(KernelCall),
    Copy,
    Depends {
        n_deps: 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.

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.

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

Affine quantization of the last axis in groups (MLX’s quantize). Outputs the packed codes (uint32, bits per value), the scales and the biases (the input’s dtype).

Fields

group_size: usize

Values per scale and bias: 32, 64 or 128.

bits: u8

Bits per value: 1–6 or 8.

Dequantize

(w, scales, biases): code · scale + bias per value, in the scales’ dtype.

Fields

group_size: usize

Values per scale and bias.

bits: u8

Bits per value.

QuantizedMatmul

(x, w, scales, biases): x · dequantize(w)ᵀ (transpose) or x · dequantize(w). w is two-dimensional, or has x’s batch dimensions.

Fields

group_size: usize

Values per scale and bias.

bits: u8

Bits per value.

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.

Implementations

impl Primitive

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

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

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