Crate nuvec_core

Expand description

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::{Array, Result};

let a = Array::from_slice(&[1.0f32, 2.0, 3.0], &[3])?;
let b = Array::full(&[3], 10.0f32)?;
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]);

Binary operations promote types (MLX’s table) and broadcast shapes (NumPy rules); views (reshape, transpose, slice, 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::{Array, Device, Stream, set_default_stream};

set_default_stream(Stream::default_for(Device::gpu(0)));
let x = Array::from_slice(&[1.0f32, 2.0, 3.0], &[3])?;
let y = &x * &x; // a Metal kernel on a Mac; an `UnsupportedDevice` error elsewhere

Status: pre-alpha. The default stream is the CPU.

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) and norm_rope (norm, RoPE and a cache write); and, through crate::softmax and crate::logsumexp over the last axis, MLX’s softmax and log-sum-exp kernels.
random: Random numbers from explicit keys or a per-thread key sequence, bit for bit as in MLX.

Structs§

Array: A lazily evaluated n-dimensional array.
Complex64: A complex number with f32 parts, the element type of DType::Complex64.
CustomCall: One call of a custom function, as recorded in the graph: its rules and the number of primal inputs. Equal (and hashed) by the identity of the function.
CustomFunction: A function with custom transform rules; see custom_function.
Device: A processor that executes array operations.
Flags: How an array’s elements are laid out in memory.
Fused: A graph of element-wise operations evaluated as one kernel, with one or more of its steps as results. Cheap to clone; equal graphs compare and hash equal.
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.
Scheduled: Arrays whose evaluation is put in order but not started (schedule).
SliceArg: One slice specification per dimension: start:stop:step with Python semantics.
Stream: An ordered queue of operations on a device.
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§

ArgReduceOp: The position of an extreme value along one axis.
BinaryOp: An element-wise function of two operands.
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.
FusedRef: An operand of a fused step: an input of the fused computation or an earlier step’s result.
Primitive: An operation recorded in the graph, with its parameters.
ReduceOp: A reduction over a set of axes.
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.
ScanOp: The combination a cumulative scan applies.
ScatterReduce: How a scatter combines an update with the value it lands on.
UnaryOp: An element-wise function of one operand.

Constants§

VERSION: The version of this crate.

Traits§

Element: A Rust type that can be stored in an array. Implemented for the element types of every DType; sealed.
ScalarOperand: A Rust number usable as a weak operand of + - * / on the right of an array.

Functions§

abs: Element-wise absolute value; the float32 magnitude for complex numbers.
add: Element-wise a + b, with promotion and broadcasting. Integers wrap; for booleans this is logical OR.
addmm: alpha · (a @ b) + beta · c, with c broadcast to the product’s shape, as MLX’s addmm. On the GPU it is one kernel (the GEMM’s epilogue), so Linear layers use it for their bias.
all: Whether every element over axes (empty = all) is non-zero.
any: Whether any element over axes (empty = all) is non-zero.
arange: Evenly spaced values start, start + step, … below stop (above it for negative steps), converted to dtype. The length is ceil((stop - start) / step), as in MLX.
argmax: Index (uint32) of the maximum along axis, or over the flattened array if None. The first occurrence wins; the first NaN wins over any number.
argmin: Index (uint32) of the minimum along axis, or over the flattened array if None. The first occurrence wins; the first NaN wins over any number.
argpartition: The indices (uint32) that partition a, as partition.
argsort: The indices (uint32) that sort a along axis, as sort.
astype: Converts the elements to dtype. Float-to-integer conversion saturates (NaN becomes 0); conversion to bool tests for non-zero; complex-to-real keeps the real part.
async_eval: Starts evaluating arrays 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.
broadcast_to: Presents a as the larger shape (NumPy broadcasting rules). Shares the data.
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.
concatenate: arrays joined along axis (MLX’s concatenate): equal shapes except along axis, with dtypes promoted to a common one.
conjugate: Element-wise complex conjugate. Real arrays are returned unchanged.
contiguous: a in row-major memory: itself if it already is, else a copy (MLX’s contiguous).
conv1d: The 1-D convolution (cross-correlation) of input ([N, L, C_in]) with weight ([C_out, K, C_in / groups]), as MLX’s conv1d: output [N, L_out, C_out] with L_out = (L + 2·padding − dilation·(K − 1) − 1) / stride + 1. Channels split into groups independent groups (groups = C_in = C_out is a depthwise convolution). Products accumulate in float32 and round once, as MLX’s kernels.
cos: Element-wise cosine, in at least float32.
cummax: The running maximum (NaN propagates), as cumsum.
cummin: The running minimum (NaN propagates), as cumsum.
cumprod: The cumulative product, as cumsum.
cumsum: The cumulative sum along axis (flattened without one): reverse scans from the end; an exclusive scan (inclusive = false) leaves each element out of its own output. Booleans sum as int32.
custom_function: Wraps f so that transforms use custom rules for it, set with CustomFunction::with_vjp, CustomFunction::with_jvp and CustomFunction::with_vmap.
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.
dequantize: Restores the values packed by quantize: code · scale + bias, in the scales’ dtype.
divide: Element-wise true division a / b. Integer operands are promoted to floating point.
equal: Element-wise a == b as a bool array.
erfinv: Element-wise inverse error function, in at least float32.
eval: Evaluates arrays and everything they depend on, and waits for them. Each node runs at most once, even across threads.
exp: Element-wise e^a, in at least float32.
expand_dims: Inserts a dimension of size 1 at axis (negative counts from the end of the result). Shares the data.
gather_mm: Matrix products of gathered batches, as MLX’s gather_mm (mixture-of-experts layers in full precision): output batch i is a[lhs_indices[i]] · b[rhs_indices[i]] over the flattened batch dimensions of a ([..., M, K]) and b ([..., K, N]), with the indices broadcast together and either defaulting to every batch. sorted_indices promises sorted rhs_indices, which lets one kernel walk each batch of b once.
gather_qmm: Quantized matrix products of gathered batches, as MLX’s gather_qmm (mixture-of-experts layers): output batch i is x[lhs_indices[i]] · dequantize(w)[rhs_indices[i]]ᵀ (or without the transpose), indexing the flattened batch dimensions of x ([..., M, K]) and of w, scales and biases. The indices (integers) broadcast together; either defaults to every batch of its operand in order. sorted_indices promises that rhs_indices is sorted (as after grouping rows by expert), which lets one kernel walk each weight batch once.
grad: Returns the gradient function of a scalar function of one array.
greater: Element-wise a > b as a bool array. Not defined for complex numbers.
greater_equal: Element-wise a >= b as a bool array. Not defined for complex numbers.
jvp: Forward mode: runs f on primals and returns its outputs together with the Jacobian-vector product, one tangent per output, for one tangent per primal.
less: Element-wise a < b as a bool array. Not defined for complex numbers.
less_equal: Element-wise a <= b as a bool array. Not defined for complex numbers.
log: Element-wise natural logarithm, in at least float32.
logaddexp: Element-wise log(exp(a) + exp(b)), computed stably as MLX does: the larger operand plus log1p(exp(smaller − larger)), or the larger one when the smaller is −∞ or the larger +∞. Integers are promoted to floating point; NaN propagates. Not defined for complex numbers.
logcumsumexp: The running log(Σ exp), as cumsum (computed in floating point).
logsumexp: log(sum(exp(a))) over axes (all axes if empty), computed stably as MLX does: the maximum is subtracted first (without gradient), and rows whose maximum is infinite return that maximum.
matmul: Matrix product of [..., M, K] and [..., K, N], broadcasting batch dimensions.
max: Maximum over axes (empty = all); NaN propagates. Zero-length axes are an error.
maximum: Element-wise maximum; NaN propagates. Not defined for complex numbers.
mean: Mean over axes (empty = all), in at least float32. Empty reductions are NaN.
min: Minimum over axes (empty = all); NaN propagates. Zero-length axes are an error.
minimum: Element-wise minimum; NaN propagates. Not defined for complex numbers.
multiply: Element-wise a * b, with promotion and broadcasting. Integers wrap; for booleans this is logical AND.
negative: Element-wise -a. Integers wrap. Not defined for booleans.
not_equal: Element-wise a != b as a bool array.
ones: An array of ones.
partition: a partitioned along axis (flattened without one): the element at position kth is the one a sort would put there, smaller ones come before it and larger ones after (MLX’s partition).
prod: Product over axes (empty = all), with the same widening as sum.
promote_types: The element type two operands are converted to before a binary operation, following MLX’s promotion table.
put_along_axis: A copy of a with values placed at indices along axis (MLX’s put_along_axis). indices broadcasts against a on the other axes, values against indices; with repeated indices, one of the values wins.
quantize: Quantizes the last axis of w (float16, bfloat16 or float32, at least 2-D, the last dimension divisible by group_size) in groups of group_size (32, 64 or 128) values with bits (1–6 or 8) bits each. Returns the packed codes (uint32, last dimension · bits / 32), the scales and the biases (w’s dtype, last dimension / group_size), equal to MLX’s bit for bit.
quantized_matmul: The matrix product of x with quantized weights, x · dequantize(w)ᵀ if transpose (MLX’s default, for weights stored as [out, in]) or x · dequantize(w), without materialising the dequantized matrix.
repeat: a with each element repeated repeats times along axis (of the flattened array without one), as MLX’s repeat.
reshape: The same elements in another shape; one dimension may be -1 (inferred). Shares the data when the layout permits, otherwise copies.
scatter_add_along_axis: A copy of a with values added at indices along axis; repeated indices accumulate (MLX’s scatter_add_axis, the gradient of take_along_axis).
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.
sigmoid: Element-wise logistic function 1 / (1 + e^-a), in at least float32.
sign: Element-wise sign: -1, 0 or 1 in a’s dtype, as MLX’s sign ((a > 0) - (a < 0), so NaN gives 0). Unsigned integers and booleans give a != 0. Complex numbers are not supported.
sin: Element-wise sine, in at least float32.
slice: Selects start:stop:step in every dimension, with Python semantics (negative indices count from the end, bounds are clamped, steps may be negative). Shares the data.
slice_update: a unchanged, but no gradient flows through it in crate::vjp, crate::jvp and A copy of a with the region selected as by slice() replaced by update (MLX’s slice_update, a[...] = update): update broadcasts to the region’s shape and takes a’s dtype.
softmax: exp(a) / sum(exp(a)) over axes (all axes if empty), with the maximum subtracted first (without gradient), as MLX does.
sort: a sorted along axis (flattened without one): ascending, stable, NaN last.
split: a split into parts equal parts along axis (MLX’s split with a count); views.
split_at: a split along axis before each of the (non-decreasing) positions indices, as MLX’s split with indices: indices.len() + 1 views.
sqrt: Element-wise square root, in at least float32.
squeeze: Removes dimensions of size 1: the given axes, or all of them if axes is empty. Shares the data.
stack: arrays (one shape) stacked along a new axis axis (MLX’s stack).
stop_gradient: crate::grad. Shares the data.
subtract: Element-wise a - b, with promotion and broadcasting. Integers wrap; for booleans this is logical XOR.
sum: Sum over axes (empty = all). Integers and bools widen to 32 or 64 bits (as in MLX); half and single precision accumulate in double precision.
take: The elements of a at indices along axis (MLX’s take): the axis is replaced by indices’ dimensions. Without an axis, a is flattened first.
take_along_axis: The elements of a at indices along axis (MLX’s take_along_axis): indices has a’s number of dimensions, and both broadcast on the other axes.
tanh: Element-wise hyperbolic tangent, in at least float32.
topk: The k largest elements of a along axis (flattened without one), in ascending order (MLX’s topk; MLX leaves the order unspecified).
transpose: Reorders the axes: output axis i is input axis axes[i]. Shares the data.
value_and_grad: Returns a function computing f’s value and its gradients with respect to the arguments at positions argnums.
vjp: Reverse mode: 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.
where_: Element-wise selection: a where condition is true, otherwise b.
zeros: An array of zeros.

Type Aliases§

Result: Result with Error as the error type.

Crate nuvec_core

Crate nuvec_core Copy item path

Modules§

Structs§

Enums§

Constants§

Traits§

Functions§

Type Aliases§

Crate nuvec_core