Crate nuvec_core 

Lazy array graph, primitives, dtypes, streams and evaluation: the core of nuvec.

Operations on Arrays record a graph and return immediately. Nothing is computed until an array is evaluated, explicitly or by reading it back:

use nuvec_core::{Result, array, ops};

let a = array![1.0, 2.0, 3.0];
let b = ops::full([3], 10.0);
let c = (&a + &b) * &a; // lazy: nothing computed yet
assert!(!c.is_evaluated());
assert_eq!(c.to_vec::<f32>()?, vec![11.0, 24.0, 39.0]);

Operations are methods (x.exp().sum(-1)), with free mirrors of MLX’s names in ops. They never return Result: an invalid call gives a failed array carrying the op, the caller’s source location and the reason (Error::Op); arrays computed from it fail with the same error, which surfaces at Array::check, Array::eval, Array::item or Array::to_vec:

use nuvec_core::{OpErrorKind, array};

let x = array![[1.0, 2.0, 3.0]];
let y = x.reshape([2, -1]).exp(); // 3 elements cannot be [2, -1]
let err = y.check().unwrap_err();
assert!(matches!(err.kind(), Some(OpErrorKind::InvalidShape(_))));
assert_eq!(err.location().unwrap().file(), file!());

Binary operations promote types (MLX’s table) and broadcast shapes (NumPy rules); views (reshape, transpose, basic indexing with Array::i, broadcast_to, expand_dims, squeeze) share data.

Reductions (sum, mean, max, argmax, …) work over any axes and layouts; matmul broadcasts batch dimensions and uses a SIMD GEMM on the CPU.

Arrays live in unified memory. Operations on the GPU stream run as Metal kernels; anything without a GPU kernel (for example float64) runs on the CPU over the same memory:

use nuvec_core::{Device, array, with_stream};

let x = array![1.0, 2.0, 3.0];
// A Metal kernel on a Mac; an `UnsupportedDevice` error elsewhere.
let y = with_stream(Device::gpu(0), || &x * &x);

Each thread’s default stream is the GPU’s when Metal is available, else the CPU’s (default_device); with_stream changes it for a scope, set_default_stream for the calling thread.

Threads may evaluate arrays that share work concurrently (each node runs once). A transform traces on the calling thread: evaluating its tracers on other threads (rayon, say) while it runs is unsupported.

Stability: no promise when this crate is depended on directly; depend on nuvec, whose re-exports are semver-tracked.

Status: pre-alpha.

Re-exports

pub use tree::LeafKind;

pub use tree::Leaves;

pub use tree::LoadMode;

pub use tree::Param;

pub use tree::Static;

pub use tree::Tree;

pub use tree::TreeDef;

pub use tree::Visit;

pub use tree::VisitMut;

pub use tree::Visitor;

pub use tree::VisitorMut;

Modules

distributed

Distributed groups and collective operations, as MLX’s mx.distributed.

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_core::prelude::*;. Operations are methods of Array; the free functions (MLX’s names) are in ops, best imported as a module (use nuvec_core::ops as mx;, then mx::zeros([2, 3], DType::F32)).

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).

MappedFile

A file’s bytes in array storage: on a Mac with a GPU, the file is memory-mapped and wrapped as GPU-visible memory without a copy, so arrays made from it are views whose pages are read when first touched (by the CPU or the GPU). Elsewhere the file is read into memory. The mapping lives as long as any array made from it.

MappedRanges

Byte ranges of a file in array storage: each range memory-mapped on its own (without a copy, on a Mac with a GPU; read elsewhere), so the bytes between them are neither mapped, wired nor counted. For loads that keep some of a file’s tensors. The file must not be modified while it is mapped (see MappedFile).

NewAxis

None (np.newaxis) in an index: a new dimension of size 1.

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.

VarOptions

Rare knobs of crate::Array::var_with and crate::Array::std_with.

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.

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.

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.

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).

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.

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.

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).

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.