Crate nuvec

Expand description

A pure-Rust array framework for Apple Silicon in the spirit of MLX: lazy arrays on unified memory and composable function transforms, with no bindings to any C/C++ library.

use nuvec::prelude::*;

let x = array![[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]];
let col = array![[10], [20]];               // broadcasts over columns
let y = (&x * 2.0 + &col).t();              // lazy, float32, a view
assert_eq!(y.shape(), [3, 2]);
assert_eq!(y.to_vec::<f32>().unwrap(), [10.0, 26.0, 12.0, 28.0, 14.0, 30.0]);

Operations are methods that chain and never return Result: an invalid call gives a failed array that remembers the operation, your source location and the reason, and the error surfaces where you read results (check()?, eval()?, item()?, to_vec()?):

use nuvec::prelude::*;

let a = array![[1.0, 2.0, 3.0]];
let bad = a.matmul(&a).exp().sum(-1);       // [1, 3] @ [1, 3]: fails, nothing panics
let err = bad.check().unwrap_err();
assert_eq!(err.op_name(), Some("matmul"));  // `err` also has the file and line
let ok = a.matmul(&a.t()).sum(..).check()?; // the boundary returns `Result`
assert_eq!(ok.item::<f32>()?, 14.0);

The free functions with MLX’s names are in ops (use nuvec::ops as mx;), so MLX code ports line by line. Gradients are lazy arrays too, so they compose:

use nuvec::{Array, Result, array, grad};

let f = |x: &Array| x.multiply(x).multiply(x).sum(..);  // sum(x³)
let x = array![1.0, 2.0];
assert_eq!(grad(f)(&x)?.to_vec::<f32>()?, [3.0, 12.0]);  // 3x²
let g2 = grad(|x: &Array| Ok(grad(f)(x)?.sum(..)))(&x)?;
assert_eq!(g2.to_vec::<f32>()?, [6.0, 12.0]);

Transforms work on trees: arrays, tuples, Vecs, maps and your own structs with #[derive(Tree)] (models with nuvec-nn’s #[derive(Module)]). A gradient has the argument’s shape, and a training step over (model, optimizer) compiles with compile_mut:

use nuvec::{Array, Result, Tree, array, value_and_grad};

#[derive(Clone, Tree)]
struct Affine { w: Array, b: Array }

let x = array![1.0, 2.0];
let loss = |p: &Affine| p.w.multiply(&x).add(&p.b).sum(..); // annotate the parameter
let (value, grads) = value_and_grad(loss)(&Affine { w: array![3.0, 4.0], b: array![1.0, 1.0] })?;
assert_eq!(value.item::<f32>()?, 13.0);
assert_eq!(grads.w.to_vec::<f32>()?, [1.0, 2.0]);              // grads: Affine
nuvec::eval(&(&value, &grads))?;                              // any tree, borrowed

What is inside:

Arrays with MLX/NumPy semantics: 14 dtypes, type promotion, broadcasting, zero-copy views, indexing, reductions, sorting, scans and matmul.
Transforms that compose: grad, value_and_grad, vjp, jvp, vmap, custom_function, checkpoint, compile() (kernel fusion) and compile_mut (stateful steps), over trees (the positional forms are *_flat).
Trees (Tree, Param, #[derive(Tree)], #[derive(Static)]): paths, structure (TreeDef) and checkpoints of any tree (save, load_into).
Metal kernels ported from MLX (GEMM, GEMV, quantized products, attention, normalisation, RoPE, sort, scan), MLX-compatible affine quantization and custom kernels from source (fast::MetalKernel).
Indexing as MLX’s x[...] (Array::i) and functional updates (Array::at).
Evaluation on the GPU by default, submitted asynchronously (async_eval, schedule), or on the CPU; memory limits and wired memory on gpu::Device.
I/O (io): safetensors and npy, memory-mapped without a copy.

The default stream of each thread is the GPU’s when Metal is available, else the CPU’s (set_default_stream changes it per thread). Several threads may evaluate arrays that share work. A transform (grad, vmap, compile, …) traces on the calling thread: evaluating its tracers on another thread inside it (rayon, say) is unsupported.

The crate root holds the types, array!, evaluation, transforms and streams; the operations are Array methods with their free-function mirrors in ops; prelude imports the everyday names; graph is the graph-level API for compilers. Neural-network layers and optimizers are in nuvec-nn, the LLM runtime in nuvec-lm.

§Determinism

For one nuvec version on one GPU generation, with the same shapes, dtypes and options:

Run to run and across processes: the same operations on the same inputs give bit-identical outputs, whatever memory the allocator hands out, whatever work is in flight, wherever command buffers end, on any thread and in any process.
Build modes: precompiled kernels and runtime-compiled ones (NUVEC_METALLIB=off) give the same bits.
Row independence: a product with at most ROW_INDEPENDENT_ROWS rows gives each row the bits of the same row’s one-row product, by default and with no mode: dense matrix products (float32, float16, bfloat16; a multi-row matrix-vector kernel that reads the weights once), quantized products (every bit width and group size; untransposed weights run one product per row), and gathered (mixture-of-experts) products, quantized or not, with sorted or unsorted indices (up to ROW_INDEPENDENT_ROWS tokens of up to 8 experts each). So are the row-wise operations of a decode step: rms_norm, add_rms_norm, rope and norm_rope with per-row offsets, element-wise operations, vector attention over equal key lengths, gated_delta and gated_delta_slots.
Batch items of quantized products: a quantized product over [B, L, K] (tiled, L above the matrix-vector limit) gives each item the bits of its product alone ([1, L, K]): the split-K partitions are counted per item, so prompts of the same length prefilled together get their solo bits.
Speculative verification (the serial rule): attention with 2 to 8 causal query rows gives each row the bits of the one-query call over its causal key prefix (fast::scaled_dot_product_attention).
Routing is pure: which kernel runs depends only on shapes, dtypes, options and the GPU generation, never on memory state (fast::plan_attention reports it).
Chunk plans: the same chunk plan gives the same bits, and resuming at a chunk boundary from a ops::copy snapshot equals the uninterrupted run.

Not promised: equality across nuvec versions, GPU generations or different chunk plans (a prompt prefilled in other chunks may round differently), and equality with MLX.

Not deterministic (the order in which GPU threads commit decides the rounding):

scatter_add_along_axis (and at(..).add) on floating-point values with repeated indices, and derivatives built on it (the reverse mode of gathers, take, take_along_axis and indexing with repeated indices): float atomics (indexing.metal). With unique indices the result is exact.
put_along_axis with repeated indices and differing values (one of the writes wins). Unique indices, or one fill value, are deterministic (as a top-k sampler’s argpartition indices and fill).
fast::MetalKernels whose source uses atomics.

§Stability

Semantic versioning covers every item reachable through nuvec:: paths, and nuvec-lm; the sub-crates re-exported here make no promise when depended on directly. Runtimes can rely on arrays and operations, the fast kernels, quantized and gathered products, custom Metal kernels, compile, random numbers, evaluation, memory limits and statistics, safetensors I/O and the attention plan. The pinned items are listed in this crate’s tests/api_surface.rs.

Status: pre-alpha: the API may change between releases.

Re-exports§

pub use nuvec_compile as compile;
pub use nuvec_cpu as cpu;
pub use nuvec_gpu as gpu;
pub use nuvec_io as io;
pub use nuvec_macros as macros;

Modules§

fast: Fused operations with dedicated GPU kernels, as MLX’s mx.fast: rms_norm, layer_norm, rope, cross_entropy, scaled_dot_product_attention and gated_delta (mlx-lm’s Gated DeltaNet kernels); nuvec’s add_rms_norm (a residual connection and its norm), norm_rope (norm, RoPE and a cache write) and paged_attention with paged_write (a block-table KV cache); and, through crate::Array::softmax and crate::Array::logsumexp over the last axis, MLX’s softmax and log-sum-exp kernels.
graph: The graph-level API: primitives, fused and custom nodes, and access to an array’s node.
ops: Array operations as free functions with MLX’s names: the mirrors of the Array methods.
prelude: The everyday names in one import: use nuvec::prelude::*;. The free functions are in ops, best imported as a module: use nuvec::ops as mx;.
random: Random numbers from explicit keys or a per-thread key sequence, bit for bit as in MLX.
tree: Trees of arrays: structured values (tuples, vectors, maps, options, user structs and enums) whose array leaves have paths and a fixed order, as MLX’s tree_flatten.

Macros§

array: An array from nested literals, as MLX’s mx.array([...]): the macro’s brackets are the outermost dimension, so array![1, 2, 3] is 1-D and array![[1.0, 2.0], [3.0, 4.0]] is 2-D (up to 4-D). Unsuffixed float literals give float32 and unsuffixed integers int32; suffixed literals keep their type (array![1u8, 2, 3] is uint8), except f64, which gives float32 like MLX’s Python floats (use Array::from for a float64 array). Ragged rows do not compile.

Structs§

Array: A lazily evaluated n-dimensional array.
At: A pending functional update made by Array::at.
Complex64: A complex number with f32 parts, the element type of DType::Complex64.
CustomFunction: A function of a tree with custom transform rules; see custom_function.
CustomFunctionFlat: A function of positional arrays with custom transform rules; see custom_function_flat.
Device: A processor that executes array operations.
Ellipsis: ... in an index: as many full slices as the other entries leave dimensions.
NewAxis: None (np.newaxis) in an index: a new dimension of size 1.
Param: A module parameter: an array and a frozen flag (trainable unless frozen).
ScanOptions: Rare knobs of the scans (crate::Array::cumsum_with and the other *_with scans).
Scheduled: Arrays whose evaluation is put in order but not started (schedule).
Stepped: A range with a step, made by Step::step: the range selects positions and the step walks them, backwards when negative (ndarray’s rule). (0..6).step(2) is Python’s 0:6:2, (..).step(-1) is ::-1, (1..4).step(-1) is 3:0:-1 (positions 3, 2, 1).
Stream: An ordered queue of operations on a device.
TreeDef: The structure of a tree: its root type, the ordered (path, kind, trainable) of every leaf, the active enum variants, the mode flags and a hash of the static fields. Equal structures have equal TreeDefs (see the module documentation).
VarOptions: Rare knobs of crate::Array::var_with and crate::Array::std_with.
Visit: The read-only visitor that Leaves::visit reports to: it builds paths and dispatches leaves, child trees, enum variants, mode flags and static fields to a Visitor.
VisitMut: The mutable visitor that Tree::visit_mut reports to.
Vmap: A function vectorized over batch axes (vmap): configure it with Vmap::in_axes/Vmap::out_axes, then call it with Vmap::call or turn it into a closure with Vmap::into_fn (to compose with compile and the other transforms).
bf16: A 16-bit floating point type implementing the bfloat16 format.
f16: A 16-bit floating point type implementing the IEEE 754-2008 standard binary16 a.k.a “half” format.

Enums§

Category: The broad kind of an element type, used for type promotion.
DType: The element type of an array.
DeviceKind: The kind of processor a device is.
Error: Why an operation failed.
LeafKind: What a leaf is (see the module documentation).
LoadMode: How Tree::load_flat treats names it does not find.
OpErrorKind: The reason an operation rejected its arguments (see Error::Op).
Scalar: A number whose dtype is decided by the array it is combined with.
ScalarOrArray: An operand of a binary operation: an array, or a weak scalar.

Constants§

ROW_INDEPENDENT_ROWS: Products with at most this many rows give each row the bits of the same row’s one-row product (the row-independence contract): dense matrix products in float32, float16 and bfloat16, quantized products (every bit width and group size), and the gathered (mixture-of-experts) products. See the nuvec crate’s “Determinism” section.
VERSION: The version of this crate.

Traits§

ArrayIndex: An index of Array::i and Array::at: an i32 or usize, a range over either (1..3, 1.., ..3, 1..=2, ..=2, with an optional .step(n)), .., Ellipsis, NewAxis, an integer index array (&Array or Array), or a tuple of up to eight of these.
Axes: Axes argument of reductions and shape ops: one axis (-1, x.ndim() - 1), a list ([0, 2], &axes[..], vec![…], or a reference to any of these), all axes (.., MLX’s axis=None) or the trailing axes from one on (1..).
Axis: The axis argument of operations that work along one axis (argmax, sort, the scans, take, split, concatenate, …): an i32 or usize axis (negative counts from the end), or .. for “the flattened array” where the operation has that meaning (as MLX’s axis=None); elsewhere .. fails with crate::OpErrorKind::InvalidAxis.
Element: A Rust type that can be stored in an array. Implemented for the element types of every DType; sealed.
GradOutput: What a function differentiated by value_and_grad returns: the loss alone (an Array) or the loss with auxiliary outputs ((Array, A) for any tree A, such as logits or metrics), which are returned as values, without gradients.
InAxes: One batch axis per top-level argument of a Vmap (None: not batched): an i32 or Option<i32> (one axis for every argument) or a tuple of them (one per argument).
InAxis: One entry of an InAxes tuple: an i32 or Option<i32>.
IntoResult: A closure’s return value: a tree T, or a Result<T>. So transform closures may return an array directly (|x: &Array| x.multiply(x).sum(..)) or a Result (with ? inside).
IntoShape: A shape argument: [2, 3] (of i32 or usize), a slice, a vector, a reference to any of these, or () for a 0-D array. Every dimension must be in 0..=i32::MAX; reshape also accepts one -1 (inferred). Anything else makes the operation fail with crate::OpErrorKind::InvalidShape.
Leaves: Anything whose array leaves can be visited in a fixed order with their paths: owned trees and borrowed views (&T, (&A, &B), [T]). Object-safe.
OutAxes: One output axis per top-level output of a Vmap: an i32 (every output) or a tuple of them.
ScalarOperand: A Rust number usable as a weak operand of + - * / on the right of an array.
Static: A value hashed into a TreeDef: the type of a #[tree(skip)] field (a configuration, a hyperparameter). A different value is a different structure, so a compiled function retraces instead of replaying a stale constant.
Step: Adds a step to a range used as an index: x.i((0..6).step(2)), x.i((..).step(-1)) (see Stepped for negative steps).
Tree: An owned tree of arrays: rebuilt with new leaves of the same structure, mutated in place.
Visitor: Receives what a Visit reports, with full paths. Every method but Visitor::leaf has an empty default.
VisitorMut: Receives what a VisitMut reports, with full paths.

Functions§

async_eval: Starts evaluating the leaves of tree (any Leaves value, as eval) and returns without waiting for GPU work: reading the arrays (or evaluating them) later waits for exactly the work they need. As in MLX, use it to overlap building the next computation with the GPU running this one.
at_least_float: d if it is floating point or complex, otherwise the promotion of d with float32. Transcendental functions and division compute in this type.
checkpoint: Returns f with gradient checkpointing: values and derivatives equal f’s, but f’s intermediates are not kept for the backward pass; f runs again to differentiate. Random numbers drawn inside f (dropout) are drawn again with the same key, so the recomputation sees the same masks.
checkpoint_flat: Returns f (a function of positional arrays) with gradient checkpointing; the flat form of checkpoint.
compile: Returns f compiled: a function with the same results that traces f once per input structure (the TreeDef: paths, leaf kinds, trainable bits, enum variants, mode flags and static fields), leaf dtypes and shapes, and default stream, and replays the simplified, fused graph on later calls.
compile_flat: compile over positional arrays: traces once per input signature (dtypes and shapes) and default stream.
compile_mut: Compiles a stateful step, as MLX’s mx.compile(step, inputs=state, outputs=state): f(&mut state, &inputs) where state is a tree (typically (model, optimizer)).
custom_function: Wraps f (a function of a tree, returning a tree) so that transforms use custom rules for it, set with CustomFunction::with_vjp, CustomFunction::with_jvp and CustomFunction::with_vmap.
custom_function_flat: Wraps f (a function of positional arrays) so that transforms use custom rules for it, set with CustomFunctionFlat::with_vjp, CustomFunctionFlat::with_jvp and CustomFunctionFlat::with_vmap; the flat form of custom_function.
default_device: The default device: the GPU when Metal is available, else the CPU (as in MLX).
default_stream: The stream new operations on this thread are recorded on: the default device’s first stream until set_default_stream is called on this thread.
eval: Evaluates the leaves of tree and everything they depend on, and waits for them. Each node runs at most once, even across threads.
grad: Returns the gradient function of f (a scalar function of a tree): the gradient with respect to the trainable leaves, as a tree of the argument’s type (see value_and_grad).
grad_flat: Returns the gradient function of f with respect to the positional arguments at argnums (the flat form of grad).
jvp: Forward mode over trees: f’s output at primals and the Jacobian-vector product for tangents (a tree like primals; tangents of leaves that are not differentiated are ignored), as a tree like the output.
jvp_flat: Forward mode over positional arrays (the flat form of jvp): runs f on primals and returns its outputs together with the Jacobian-vector product, one tangent per output, for one tangent per primal.
load_into: Reads a file written by save into tree by path, strictly: every leaf must be in the file with its shape, and nothing else (loaded dtypes are kept). On any problem the error lists every one and tree is unchanged.
promote_types: The element type two operands are converted to before a binary operation, following MLX’s promotion table.
save: Writes the leaves of tree as a safetensors file, each named by its path (MLX’s tree_flatten names), evaluating them.
schedule: arrays with the work they need put in order, to be started later by Scheduled::async_eval or Scheduled::eval with less delay: ordering a large graph takes a while, which can then overlap other work, such as the GPU running the previous step while the next one waits for its result.
set_default_stream: Sets the stream new operations on this thread are recorded on.
value_and_grad: Returns a function computing f’s value and its gradient with respect to the trainable leaves of its argument, as a tree of the argument’s type (grads.layers[0].weight).
value_and_grad_flat: Returns a function computing f’s value and its gradients with respect to the positional arguments at argnums (the flat form of value_and_grad).
vjp: Reverse mode over trees: f’s output at primals and the vector-Jacobian product for cotangents (a tree like the output: same paths, shapes and dtypes, else an error naming the path), as a tree like primals (zeros for leaves that are not differentiated).
vjp_flat: Reverse mode over positional arrays (the flat form of vjp): runs f on primals and returns its outputs together with the vector-Jacobian product, one gradient per primal, for one cotangent per output.
vmap: Returns f vectorized over a batch axis of its inputs (axis 0 of every top-level argument and output by default; see Vmap::in_axes and Vmap::out_axes).
vmap_flat: Returns f vectorized over a batch axis, over positional arrays (the flat form of vmap).
with_stream: Runs f with stream as this thread’s default stream, and restores the previous default afterwards, also when f panics. A device stands for its first stream.

Type Aliases§

Result: Result with Error as the error type.

Derive Macros§

Static: Derives Static (the type of a hashed #[tree(skip)] field) by hashing every field in order (enums: the variant index, then its fields).
Tree: Derives Leaves and Tree for a struct (named, tuple or unit) or an enum: the type’s array leaves get paths from its field names, as MLX’s tree_flatten.

Crate nuvec

Crate nuvec Copy item path

§Determinism

§Stability

Re-exports§

Modules§

Macros§

Structs§

Enums§

Constants§

Traits§

Functions§

Type Aliases§

Derive Macros§

Crate nuvec