Skip to main content

Crate nuvec

Crate nuvec 

Source
Expand description

A pure-Rust array framework for Apple Silicon in the spirit of MLX: lazy n-dimensional arrays (tensors) 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 in memory, device facts in metal.
  • 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: dense matrix products (float32, float16, bfloat16; a multi-row matrix-vector kernel that reads the weights once) always, and quantized and gathered (mixture-of-experts) products under with_row_exact (every bit width and group size, sorted or unsorted indices, up to ROW_INDEPENDENT_ROWS tokens of up to 8 experts each). Outside it, quantized and gathered products route as MLX’s do: from 2 rows, multi-vector kernels that read the weights once give a row other low bits than its one-row product. The row-wise operations of a decode step are row-independent always: 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: under with_row_exact, 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. Outside it, as in MLX, K is split only for one item.
  • 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); in float32 at head sizes 192 and 256, for at most 32 / gqa rows (more compose). With speculative::accept_greedy and fast::gated_delta_steps_with, run under with_row_exact, a verification step decodes exactly the tokens serial decoding would.
  • Sliding windows: a vector-kernel attention row with a window gives the bits of the unmasked one-query call over exactly its window’s keys, so a compacted contiguous cache and the paged cache give the same bits (a batch-invariant Batcher serves window models’ requests their solo tokens).
  • Routing is pure: which kernel runs depends only on shapes, dtypes, options, the row-exact mode 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 make no promise when depended on directly, and the backend crates are not reachable from here. The public modules are ops, fast, random, graph, tree, speculative, row_exact, compile, distributed, io, memory, metal, macros and prelude. 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, device facts and captures, 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_io as io;
pub use nuvec_macros as macros;

Modules§

distributed
Distributed groups and collectives (MLX’s mx.distributed): all_sum, all_gather, sum_scatter, send, … over a Group, on MLX’s TCP ring or (feature jaccl, experimental) MLX’s JACCL over Thunderbolt 5 RDMA; init reads MLX’s launch environment.
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.
memory
Memory statistics and limits (MLX’s mx.get_active_memory, mx.set_memory_limit, …), in bytes, of the shared Metal device’s allocator. Without a Metal device the statistics are zero and setting a limit returns 0 and changes nothing, as MLX’s CPU build.
metal
The Metal device (MLX’s mx.metal): availability, device facts, the Neural Accelerator kernels’ status, GPU captures and the kernel profile. Every function works without a Metal device.
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.
speculative
Acceptance for speculative decoding, on the device: from a verification’s logits (or probabilities) and the draft, one [3] uint32 array [accepted, next, eos_hit], so a verification step costs one small read.
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.
GatherQmm
Options of gather_qmm, built with setters. GatherQmm::default() multiplies by the transposed weights (MLX’s default).
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).
Quant
How weights are quantized: a mode, a group size and a bit width. Quant::new gives MLX’s defaults for a mode; an unsupported combination fails in the operation that uses it (a failed array, as every operation’s errors).
Quantized
Quantized weights: packed codes (uint32), one scale per group (affine: the weights’ dtype; floating-point modes: uint8 codes), affine biases, and how they were quantized.
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).
QuantMode
A quantization mode (MLX’s mode strings: "affine", "mxfp4", "mxfp8", "nvfp4").
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: dense matrix products in float32, float16 and bfloat16 always, and quantized and gathered (mixture-of-experts) products evaluated under with_row_exact. 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. For a rank known only at run time, a slice or Vec of ranges indexes the leading axes, one range each (&[1..3, 0..2]); integers are not accepted there, so a list is never mistaken for a gather.
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.
row_exact
Whether products evaluated or built on this thread are row-exact (with_row_exact).
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_row_exact
Runs f with row-exact products on this thread: quantized and gathered products of up to ROW_INDEPENDENT_ROWS rows (tokens, for gathers) evaluated or built inside give each row the bits of its one-row product, through one matrix-vector product per row (or nuvec’s few-row kernels with the same arithmetic). Outside it they route as MLX does, whose multi-vector kernels are faster from 2 rows on but give a row other bits than its one-row product. Speculative verification runs in this mode, and so can batched serving that must equal each request run alone.
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.