Struct Param 

pub struct Param { /* private fields */ }

A module parameter: an array and a frozen flag (trainable unless frozen).

It reads as an Array (Deref), so methods, &Param arguments where &Array is expected, operators (&x * &p, &p + 1.0) and free functions (ops::add(&p, &x)) work as on the array. Gradients, Tree::map and clones keep the flag.

Implementations

impl Param

pub fn new(value: Array) -> Self

A trainable parameter.

pub fn frozen(value: Array) -> Self

A parameter frozen at construction (a LoRA base, say).

pub fn is_frozen(&self) -> bool

Whether the parameter is frozen (not trained).

pub fn set_frozen(&mut self, frozen: bool)

Freezes (true) or unfreezes (false) the parameter.

pub fn value(&self) -> &Array

The array.

pub fn value_mut(&mut self) -> &mut Array

The array, mutably (the frozen flag is kept).

pub fn into_inner(self) -> Array

The array, consuming the parameter.

Methods from Deref<Target = Array>

pub fn is_failed(&self) -> bool

Whether this is a failed array: the operation that produced it rejected its arguments, or one of its inputs had failed, when it was recorded. O(1); errors that happen during evaluation are not visible here (they surface from Array::eval).

pub fn error(&self) -> Option<&Error>

The error of a failed array (see Array::is_failed), without evaluating.

pub fn shape(&self) -> &[i32]

Dimension sizes (MLX’s Shape: i32).

A failed array keeps a best-effort shape: the one the operation would have produced if that can be inferred, else its first operand’s. Call Array::check first when host logic depends on the shape.

pub fn shape_n<const N: usize>(&self) -> Result<[i32; N]>

The shape as an array of N sizes, for destructuring and host decisions (let [b, l, d] = x.shape_n::<3>()?;): the array’s error if it is a failed array (whose shape is only a best guess), or an error if it does not have N dimensions.

use nuvec_core::array;

let x = array![[1, 2, 3], [4, 5, 6]];
let [rows, cols] = x.shape_n::<2>().unwrap();
assert_eq!((rows, cols), (2, 3));
assert!(x.shape_n::<3>().is_err());

pub fn ndim(&self) -> usize

Number of dimensions.

pub fn size(&self) -> usize

Number of elements.

pub fn dtype(&self) -> DType

Element type.

pub fn itemsize(&self) -> usize

Size of one element in bytes.

pub fn nbytes(&self) -> usize

Size of all elements in bytes.

pub fn stream(&self) -> Stream

The stream the producing operation was recorded on.

pub fn device(&self) -> Device

The device the producing operation runs on.

pub fn is_evaluated(&self) -> bool

Whether the array has been evaluated (successfully or not).

pub fn eval(&self) -> Result<()>

Evaluates this array, if it has not been evaluated yet.

pub fn to_vec<T: Element>(&self) -> Result<Vec<T>>

The elements in row-major order, evaluating first if needed.

pub fn item<T: Element>(&self) -> Result<T>

The single element of a one-element array, evaluating first if needed.

pub fn buffer_id(&self) -> Option<usize>

The identity of the memory the evaluated array reads, or None before evaluation (and for a failed array). Arrays with equal ids share storage (views of one array, or tensors of one mapped or read checkpoint file); an array’s id stays the same while it is alive. The analogue of comparing MLX’s data pointers for aliasing; it is not an address to read from.

use nuvec_core::{array, Array};

let a = array![[1, 2], [3, 4]].add(1);
let row = a.i(1);
assert_eq!(a.buffer_id(), None); // not evaluated yet
a.eval().unwrap();
row.eval().unwrap();
assert!(a.buffer_id().is_some());
assert_eq!(row.buffer_id(), a.buffer_id()); // a view shares the storage
assert_ne!(Array::from(1).buffer_id(), a.buffer_id());

pub fn i(&self, index: impl ArrayIndex) -> Array

The array indexed as MLX’s x[...]: integers, ranges (1..3, ..=2, (0..6).step(2), (..).step(-1)), .., Ellipsis, NewAxis, integer index arrays, and tuples of these (one entry per dimension, leading dimensions first).

Basic indices give views (no copy); integer arrays gather. Integers out of bounds, two ellipses, too many entries or a zero step make the result a failed array (crate::OpErrorKind::InvalidIndex); slices clamp as in Python.

use nuvec_core::{Array, Ellipsis, NewAxis, Step};

let x = Array::from_slice(&(0..24).collect::<Vec<i32>>(), [2, 3, 4]);
assert_eq!(x.i(0).shape(), [3, 4]);                          // x[0]
assert_eq!(x.i((.., 1..3)).shape(), [2, 2, 4]);              // x[:, 1:3]
assert_eq!(x.i((Ellipsis, -1)).to_vec::<i32>().unwrap()[0], 3); // x[..., -1]
assert_eq!(x.i((.., NewAxis)).shape(), [2, 1, 3, 4]);        // x[:, None]
assert_eq!(x.i((.., .., (0..4).step(2))).shape(), [2, 3, 2]); // x[:, :, 0:4:2]
let rows = Array::from_slice(&[1i32, 0], [2]);
assert_eq!(x.i((0, &rows)).shape(), [2, 4]);                 // x[0, rows]
assert!(x.i(5).is_failed());

pub fn slice_update(&self, update: &Array, index: impl ArrayIndex) -> Array

A copy of this array with the region selected by index replaced by update (MLX’s slice_update, a[index] = update): the borrowing form of Array::at with At::set, which writes in place only when the base is not shared.

pub fn add(&self, other: impl Into<ScalarOrArray>) -> Array

Element-wise self + other. Integers wrap; for booleans this is logical OR.

use nuvec_core::array;

let x = array![1.0, 2.0];
assert_eq!(x.add(1.0).to_vec::<f32>().unwrap(), [2.0, 3.0]);
assert_eq!(x.add(&x).to_vec::<f32>().unwrap(), [2.0, 4.0]);

pub fn subtract(&self, other: impl Into<ScalarOrArray>) -> Array

Element-wise self - other. Integers wrap; for booleans this is logical XOR.

use nuvec_core::array;

assert_eq!(array![3, 5].subtract(1).to_vec::<i32>().unwrap(), [2, 4]);

pub fn multiply(&self, other: impl Into<ScalarOrArray>) -> Array

Element-wise self * other. Integers wrap; for booleans this is logical AND.

use nuvec_core::{DType, array};

let y = array![1u8, 2].multiply(3);
assert_eq!((y.dtype(), y.to_vec::<u8>().unwrap()), (DType::U8, vec![3, 6]));

pub fn divide(&self, other: impl Into<ScalarOrArray>) -> Array

Element-wise true division; integers are promoted to floating point.

use nuvec_core::array;

assert_eq!(array![1, 3].divide(2).to_vec::<f32>().unwrap(), [0.5, 1.5]);

pub fn maximum(&self, other: impl Into<ScalarOrArray>) -> Array

Element-wise maximum; NaN propagates. Not defined for complex numbers.

use nuvec_core::array;

assert_eq!(array![-1.0, 2.0].maximum(0.0).to_vec::<f32>().unwrap(), [0.0, 2.0]);

pub fn minimum(&self, other: impl Into<ScalarOrArray>) -> Array

Element-wise minimum; NaN propagates. Not defined for complex numbers.

pub fn logaddexp(&self, other: impl Into<ScalarOrArray>) -> Array

Element-wise log(exp(self) + exp(other)), computed stably, in floating point.

pub fn floor_divide(&self, other: impl Into<ScalarOrArray>) -> Array

Element-wise floor(self / other) (MLX’s floor_divide, Python’s //): for floats the floor of the quotient, for integers the quotient rounded toward negative infinity, in the promoted dtype. An integer division by zero gives 0 (MLX leaves it undefined). Not defined for bool or complex; no derivative.

use nuvec_core::array;

assert_eq!(array![7, -7, 7].floor_divide(array![2, 2, -2]).to_vec::<i32>().unwrap(), [3, -4, -4]);
assert_eq!(array![7.5f32, -7.5].floor_divide(2.0).to_vec::<f32>().unwrap(), [3.0, -4.0]);

pub fn power(&self, other: impl Into<ScalarOrArray>) -> Array

Element-wise self to the power other (MLX’s power, Python’s **), in the promoted dtype: pow for floats, repeated squaring for integers (wrapping; 0 for a negative exponent, which MLX leaves undefined). Not defined for bool or complex. The derivative in other is taken as 0 where self ≤ 0.

use nuvec_core::array;

assert_eq!(array![2, 3, 2].power(array![3, 2, -1]).to_vec::<i32>().unwrap(), [8, 9, 0]);
assert_eq!(array![4.0f32, 2.0].power(array![0.5f32, 3.0]).to_vec::<f32>().unwrap(), [2.0, 8.0]);

pub fn equal(&self, other: impl Into<ScalarOrArray>) -> Array

Element-wise self == other as a bool array.

use nuvec_core::array;

assert_eq!(array![1, 2].equal(2).to_vec::<bool>().unwrap(), [false, true]);

pub fn not_equal(&self, other: impl Into<ScalarOrArray>) -> Array

Element-wise self != other as a bool array.

pub fn less(&self, other: impl Into<ScalarOrArray>) -> Array

Element-wise self < other as a bool array.

pub fn less_equal(&self, other: impl Into<ScalarOrArray>) -> Array

Element-wise self <= other as a bool array.

pub fn greater(&self, other: impl Into<ScalarOrArray>) -> Array

Element-wise self > other as a bool array.

use nuvec_core::array;

assert_eq!(array![1.0, 3.0].greater(2.0).to_vec::<bool>().unwrap(), [false, true]);

pub fn greater_equal(&self, other: impl Into<ScalarOrArray>) -> Array

Element-wise self >= other as a bool array.

pub fn negative(&self) -> Array

Element-wise -self. Integers wrap. Not defined for booleans.

pub fn abs(&self) -> Array

Element-wise absolute value; the float32 magnitude for complex numbers.

use nuvec_core::array;

assert_eq!(array![-2, 3].abs().to_vec::<i32>().unwrap(), [2, 3]);

pub fn exp(&self) -> Array

Element-wise e^self, in at least float32.

use nuvec_core::array;

assert_eq!(array![0.0].exp().item::<f32>().unwrap(), 1.0);

pub fn log(&self) -> Array

Element-wise natural logarithm, in at least float32.

use nuvec_core::array;

assert_eq!(array![1.0].log().item::<f32>().unwrap(), 0.0);

pub fn sqrt(&self) -> Array

Element-wise square root, in at least float32.

use nuvec_core::array;

assert_eq!(array![4, 9].sqrt().to_vec::<f32>().unwrap(), [2.0, 3.0]);

pub fn sin(&self) -> Array

Element-wise sine, in at least float32.

pub fn cos(&self) -> Array

Element-wise cosine, in at least float32.

pub fn tanh(&self) -> Array

Element-wise hyperbolic tangent, in at least float32.

use nuvec_core::array;

assert_eq!(array![0.0].tanh().item::<f32>().unwrap(), 0.0);

pub fn sigmoid(&self) -> Array

Element-wise logistic function 1 / (1 + e^-self), in at least float32.

use nuvec_core::array;

assert_eq!(array![0.0].sigmoid().item::<f32>().unwrap(), 0.5);

pub fn erfinv(&self) -> Array

Element-wise inverse error function, in at least float32.

pub fn sign(&self) -> Array

Element-wise sign: -1, 0 or 1 in the array’s dtype (NaN gives 0).

pub fn conjugate(&self) -> Array

Element-wise complex conjugate; real arrays are returned unchanged.

pub fn stop_gradient(&self) -> Array

The array unchanged, but no gradient flows through it. Shares the data.

use nuvec_core::{Array, array, grad};

let g = grad(|x: &Array| x.multiply(&x.stop_gradient()).sum(..))(&array![3.0]).unwrap();
assert_eq!(g.item::<f32>().unwrap(), 3.0); // d(x·c)/dx with c = x held constant

pub fn contiguous(&self) -> Array

The array in row-major memory: itself if it already is, else a copy.

pub fn copy(&self) -> Array

An independent copy in a buffer of its own (MLX’s mx.copy), for snapshots: no compile simplification, constant folding or in-place update removes it or makes it share memory with self, so updating self in place afterwards leaves the copy unchanged. Derivatives pass through unchanged.

use nuvec_core::array;

let cache = array![1.0, 2.0, 3.0];
let snapshot = cache.copy();
let cache = cache.at(0).set(9.0); // in place where `cache` owns its buffer
assert_eq!(snapshot.to_vec::<f32>().unwrap(), [1.0, 2.0, 3.0]);
assert_eq!(cache.to_vec::<f32>().unwrap(), [9.0, 2.0, 3.0]);

Memory

The copy’s own bytes; no scratch.

pub fn select( &self, a: impl Into<ScalarOrArray>, b: impl Into<ScalarOrArray>, ) -> Array

a where self (the condition, converted to bool) is true, else b (MLX’s where); a and b may be weak Rust numbers.

use nuvec_core::array;

let m = array![true, false];
let a = array![1.0, 2.0];
assert_eq!(m.select(&a, 0.0).to_vec::<f32>().unwrap(), [1.0, 0.0]);

pub fn astype(&self, dtype: DType) -> Array

The elements converted to dtype. Float-to-integer conversion saturates (NaN becomes 0); conversion to bool tests for non-zero; complex-to-real keeps the real part.

use nuvec_core::{DType, array};

assert_eq!(array![1.7, -2.2].astype(DType::I32).to_vec::<i32>().unwrap(), [1, -2]);

pub fn matmul(&self, other: &Array) -> Array

The matrix product self @ other ([..., M, K] @ [..., K, N], batch dimensions broadcast; 1-D operands as in NumPy). The promoted type must be floating point.

use nuvec_core::array;

let a = array![[1.0, 2.0], [3.0, 4.0]];
assert_eq!(a.matmul(&a.t()).to_vec::<f32>().unwrap(), [5.0, 11.0, 11.0, 25.0]);
let bad = a.matmul(&array![1.0, 2.0, 3.0]);
assert!(bad.is_failed());

With at most ROW_INDEPENDENT_ROWS rows, each row of the result has the bits of that row’s product alone (multi-row matrix-vector kernels on the GPU).

Memory

No scratch, except split-K GEMM on the GPU (few output tiles over a long K): partitions × M × N × 4 bytes of float32 partial products, partitions being K / 16 / tiles rounded up to a power of two in 2..=32 (tiles = ⌈M/32⌉·⌈N/32⌉). Operands that are neither row- nor column-major are copied first.

pub fn addmm(&self, a: &Array, b: &Array, alpha: f32, beta: f32) -> Array

alpha · (a @ b) + beta · self, with self broadcast to the product’s shape (MLX’s addmm; one GEMM with an epilogue on the GPU).

use nuvec_core::array;

let bias = array![1.0, 1.0];
let y = bias.addmm(&array![[1.0, 2.0]], &array![[1.0, 0.0], [0.0, 1.0]], 1.0, 1.0);
assert_eq!(y.to_vec::<f32>().unwrap(), [2.0, 3.0]);

pub fn sum(&self, axes: impl Axes) -> Array

The sum over axes. Integers and bools widen to 32 or 64 bits, as in MLX.

use nuvec_core::array;

let x = array![[1, 2, 3], [4, 5, 6]];
assert_eq!(x.sum(-1).to_vec::<i32>().unwrap(), [6, 15]);
assert_eq!(x.sum(..).item::<i32>().unwrap(), 21);

pub fn sum_keepdims(&self, axes: impl Axes) -> Array

Array::sum keeping the reduced axes as size-1 dimensions.

use nuvec_core::array;

assert_eq!(array![[1, 2, 3], [4, 5, 6]].sum_keepdims(0).shape(), [1, 3]);

pub fn prod(&self, axes: impl Axes) -> Array

The product over axes, widened as Array::sum.

use nuvec_core::array;

assert_eq!(array![2, 3, 4].prod(..).item::<i32>().unwrap(), 24);

pub fn prod_keepdims(&self, axes: impl Axes) -> Array

Array::prod keeping the reduced axes.

pub fn max(&self, axes: impl Axes) -> Array

The maximum over axes; NaN propagates; zero-length axes fail.

use nuvec_core::array;

assert_eq!(array![[1, 5], [7, 2]].max(1).to_vec::<i32>().unwrap(), [5, 7]);

pub fn max_keepdims(&self, axes: impl Axes) -> Array

Array::max keeping the reduced axes.

pub fn min(&self, axes: impl Axes) -> Array

The minimum over axes; NaN propagates; zero-length axes fail.

pub fn min_keepdims(&self, axes: impl Axes) -> Array

Array::min keeping the reduced axes.

pub fn all(&self, axes: impl Axes) -> Array

Whether every element over axes is non-zero.

use nuvec_core::array;

assert!(!array![true, false].all(..).item::<bool>().unwrap());

pub fn all_keepdims(&self, axes: impl Axes) -> Array

Array::all keeping the reduced axes.

pub fn any(&self, axes: impl Axes) -> Array

Whether any element over axes is non-zero.

pub fn any_keepdims(&self, axes: impl Axes) -> Array

Array::any keeping the reduced axes.

pub fn mean(&self, axes: impl Axes) -> Array

The mean over axes, in at least float32. Empty reductions are NaN.

use nuvec_core::array;

assert_eq!(array![1, 2, 3, 4].mean(..).item::<f32>().unwrap(), 2.5);

pub fn mean_keepdims(&self, axes: impl Axes) -> Array

Array::mean keeping the reduced axes.

pub fn logsumexp(&self, axes: impl Axes) -> Array

log(sum(exp(self))) over axes, computed stably (the maximum subtracted first).

use nuvec_core::array;

let l = array![0.0, 0.0].logsumexp(..).item::<f32>().unwrap();
assert!((l - 2f32.ln()).abs() < 1e-6);

pub fn logsumexp_keepdims(&self, axes: impl Axes) -> Array

Array::logsumexp keeping the reduced axes.

pub fn var(&self, axes: impl Axes) -> Array

The variance over axes (mean((x − mean(x))²)), in at least float32.

use nuvec_core::{VarOptions, array};

let x = array![1.0, 2.0, 3.0, 4.0];
assert_eq!(x.var(..).item::<f32>().unwrap(), 1.25);
let sample = x.var_with(.., VarOptions { ddof: 1, ..Default::default() });
assert!((sample.item::<f32>().unwrap() - 5.0 / 3.0).abs() < 1e-6);

pub fn var_keepdims(&self, axes: impl Axes) -> Array

Array::var keeping the reduced axes.

pub fn var_with(&self, axes: impl Axes, options: VarOptions) -> Array

Array::var with VarOptions (ddof, keepdims).

pub fn std(&self, axes: impl Axes) -> Array

The standard deviation over axes (the square root of Array::var).

use nuvec_core::array;

assert_eq!(array![1.0, 3.0].std(..).item::<f32>().unwrap(), 1.0);

pub fn std_keepdims(&self, axes: impl Axes) -> Array

Array::std keeping the reduced axes.

pub fn std_with(&self, axes: impl Axes, options: VarOptions) -> Array

Array::std with VarOptions (ddof, keepdims).

pub fn argmax(&self, axis: impl Axis) -> Array

The index (uint32) of the maximum along one Axis (..: of the flattened array). The first occurrence wins; the first NaN wins over any number.

use nuvec_core::array;

let x = array![[1, 9, 3], [7, 2, 8]];
assert_eq!(x.argmax(-1).to_vec::<u32>().unwrap(), [1, 2]);
assert_eq!(x.argmax(..).item::<u32>().unwrap(), 1);

pub fn argmax_keepdims(&self, axis: impl Axis) -> Array

Array::argmax keeping the reduced axis (every axis for ..).

pub fn argmin(&self, axis: impl Axis) -> Array

The index (uint32) of the minimum along one axis (..: of the flattened array).

pub fn argmin_keepdims(&self, axis: impl Axis) -> Array

Array::argmin keeping the reduced axis (every axis for ..).

pub fn softmax(&self, axes: impl Axes) -> Array

exp(self) / sum(exp(self)) over axes, with the maximum subtracted first. Always MLX’s softmax(..., precise=True): float16 and bfloat16 are computed in float32 and rounded once to the input dtype, on every device and axis, so there is no precise flag and no need to cast first.

use nuvec_core::array;

assert_eq!(array![1.0, 1.0].softmax(-1).to_vec::<f32>().unwrap(), [0.5, 0.5]);

pub fn cumsum(&self, axis: impl Axis) -> Array

The cumulative sum along one axis. Booleans sum as int32.

use nuvec_core::{ScanOptions, array};

let a = array![1, 2, 3, 4];
assert_eq!(a.cumsum(0).to_vec::<i32>().unwrap(), [1, 3, 6, 10]);
let opts = ScanOptions { reverse: true, exclusive: true };
assert_eq!(a.cumsum_with(0, opts).to_vec::<i32>().unwrap(), [9, 7, 4, 0]);

pub fn cumsum_with(&self, axis: impl Axis, options: ScanOptions) -> Array

Array::cumsum with ScanOptions (reverse, exclusive).

pub fn cumprod(&self, axis: impl Axis) -> Array

The cumulative product along one axis.

use nuvec_core::array;

assert_eq!(array![1, 2, 3].cumprod(..).to_vec::<i32>().unwrap(), [1, 2, 6]);

pub fn cumprod_with(&self, axis: impl Axis, options: ScanOptions) -> Array

Array::cumprod with ScanOptions.

pub fn cummax(&self, axis: impl Axis) -> Array

The running maximum along one axis (NaN propagates).

pub fn cummax_with(&self, axis: impl Axis, options: ScanOptions) -> Array

Array::cummax with ScanOptions.

pub fn cummin(&self, axis: impl Axis) -> Array

The running minimum along one axis (NaN propagates).

pub fn cummin_with(&self, axis: impl Axis, options: ScanOptions) -> Array

Array::cummin with ScanOptions.

pub fn logcumsumexp(&self, axis: impl Axis) -> Array

The running log(Σ exp) along one axis, in floating point.

pub fn logcumsumexp_with(&self, axis: impl Axis, options: ScanOptions) -> Array

Array::logcumsumexp with ScanOptions.

pub fn sort(&self, axis: impl Axis) -> Array

The array sorted along one axis: ascending, stable, NaN last.

use nuvec_core::array;

assert_eq!(array![3, 1, 2].sort(-1).to_vec::<i32>().unwrap(), [1, 2, 3]);

pub fn argsort(&self, axis: impl Axis) -> Array

The indices (uint32) that sort the array along one axis, as Array::sort.

use nuvec_core::array;

assert_eq!(array![3, 1, 2].argsort(0).to_vec::<u32>().unwrap(), [1, 2, 0]);

pub fn partition(&self, kth: i32, axis: impl Axis) -> Array

The array partitioned along one axis: the element at kth (negative counts from the end) is the one a sort would put there, smaller ones before, larger ones after.

pub fn argpartition(&self, kth: i32, axis: impl Axis) -> Array

The indices (uint32) that partition the array, as Array::partition.

pub fn topk(&self, k: i32, axis: impl Axis) -> Array

The k largest elements along one axis, in ascending order.

use nuvec_core::array;

assert_eq!(array![5, 1, 9, 3, 7].topk(2, 0).to_vec::<i32>().unwrap(), [7, 9]);

pub fn reshape(&self, shape: impl IntoShape) -> Array

The same elements in another shape (IntoShape: [b, l, -1], one -1 inferred). Shares the data when the layout permits.

use nuvec_core::array;

let x = array![[1, 2, 3], [4, 5, 6]];
let [b, n] = x.shape_n::<2>().unwrap();
assert_eq!(x.reshape([n, b]).shape(), [3, 2]);
assert_eq!(x.reshape([-1]).shape(), [6]);

pub fn transpose(&self, axes: impl Axes) -> Array

The axes reordered: output axis i is input axis axes[i] (.. reverses them, as MLX’s default). Shares the data.

use nuvec_core::{DType, ops::zeros};

let x = zeros([2, 3, 4], DType::F32);
assert_eq!(x.transpose([0, 2, 1]).shape(), [2, 4, 3]);
assert_eq!(x.transpose(..).shape(), [4, 3, 2]);

pub fn broadcast_to(&self, shape: impl IntoShape) -> Array

The array presented as the larger shape (NumPy broadcasting). Shares the data.

use nuvec_core::array;

assert_eq!(array![1, 2].broadcast_to([3, 2]).shape(), [3, 2]);

pub fn expand_dims(&self, axes: impl Axes) -> Array

Size-1 dimensions inserted at axes (positions in the result). Shares the data.

use nuvec_core::array;

assert_eq!(array![1, 2].expand_dims([0, -1]).shape(), [1, 2, 1]);

pub fn squeeze(&self, axes: impl Axes) -> Array

Size-1 dimensions removed: the given axes, or all of them for ... Shares the data.

use nuvec_core::array;

assert_eq!(array![[[1], [2]]].squeeze(..).shape(), [2]);

pub fn repeat(&self, repeats: i32, axis: impl Axis) -> Array

Each element repeated repeats times along one axis (..: the flattened array).

use nuvec_core::array;

assert_eq!(array![1, 2].repeat(2, 0).to_vec::<i32>().unwrap(), [1, 1, 2, 2]);

pub fn split(&self, parts: i32, axis: impl Axis) -> Vec<Array>

The array split into parts equal parts along axis (views); parts failed arrays if the axis does not divide.

use nuvec_core::array;

let parts = array![1, 2, 3, 4, 5, 6].split(3, 0);
assert_eq!(parts[2].to_vec::<i32>().unwrap(), [5, 6]);

pub fn split_at(&self, indices: &[i32], axis: impl Axis) -> Vec<Array>

The array split before each of the (non-decreasing) positions indices along axis: indices.len() + 1 views.

use nuvec_core::array;

let parts = array![1, 2, 3, 4].split_at(&[1, 3], 0);
assert_eq!(parts[1].to_vec::<i32>().unwrap(), [2, 3]);

pub fn take(&self, indices: &Array, axis: impl Axis) -> Array

The elements at indices along one axis (..: the flattened array): the axis is replaced by indices’ dimensions (MLX’s take). Negative indices count from the end.

use nuvec_core::array;

let table = array![[0.0, 0.5], [1.0, 1.5], [2.0, 2.5]];
let rows = table.take(&array![2, 0], 0);
assert_eq!(rows.to_vec::<f32>().unwrap(), [2.0, 2.5, 0.0, 0.5]);

pub fn take_along_axis(&self, indices: &Array, axis: impl Axis) -> Array

The elements at indices along axis (MLX’s take_along_axis).

use nuvec_core::array;

let a = array![[10.0, 20.0], [30.0, 40.0]];
let i = array![[1], [0]];
assert_eq!(a.take_along_axis(&i, 1).to_vec::<f32>().unwrap(), [20.0, 30.0]);

pub fn put_along_axis( &self, indices: &Array, values: impl Into<ScalarOrArray>, axis: impl Axis, ) -> Array

A copy with values placed at indices along axis (MLX’s put_along_axis).

use nuvec_core::array;

let a = array![0, 0, 0, 0];
let b = a.put_along_axis(&array![3, -4], &array![7, 9], 0);
assert_eq!(b.to_vec::<i32>().unwrap(), [9, 0, 0, 7]);

pub fn scatter_add_along_axis( &self, indices: &Array, values: impl Into<ScalarOrArray>, axis: impl Axis, ) -> Array

A copy with values added at indices along axis; repeated indices accumulate.

use nuvec_core::array;

let a = array![0, 0];
assert_eq!(a.scatter_add_along_axis(&array![1, 1], 5, 0).to_vec::<i32>().unwrap(), [0, 10]);

pub fn quantize(&self, quant: impl Into<Quant>) -> Quantized

The matrix quantized along its last axis (MLX’s quantize): affine with Quant::affine(group_size, bits), or a floating-point mode (QuantMode::Mxfp4, Mxfp8, Nvfp4, whose codes and scales equal MLX’s Metal quantizer’s).

use nuvec_core::{DType, Quant, QuantMode, ops::ones};

let q = ones([2, 64], DType::F32).quantize(Quant::affine(64, 4));
assert_eq!((q.weight.shape(), q.scales.shape()), (&[2, 8][..], &[2, 1][..]));
assert_eq!(q.biases.unwrap().shape(), [2, 1]);

// MXFP4: E2M1 codes, one uint8 E8M0 scale per 32 values, no biases.
let q = ones([2, 64], DType::F32).quantize(QuantMode::Mxfp4);
assert_eq!((q.weight.shape(), q.scales.shape()), (&[2, 8][..], &[2, 2][..]));
assert_eq!((q.scales.dtype(), q.biases.is_none()), (DType::U8, true));
let y = ones([1, 64], DType::F32).quantized_matmul(&q, true);
assert_eq!(y.to_vec::<f32>().unwrap(), [64.0, 64.0]);

pub fn quantized_matmul(&self, w: &Quantized, transpose: bool) -> Array

self · dequantize(w)ᵀ (transpose) or self · dequantize(w). Up to ROW_INDEPENDENT_ROWS rows run as matrix-vector products, each row with the bits of its own product. Affine promotes self and the scales to one dtype; the floating-point modes keep self’s dtype.

Memory

No scratch, except quantized split-K (many rows, transposed weights, one batch): split × M × N × 4 bytes, split = min(512 / tiles, K / max(group, 32)) reduced until it divides K (tiles = ⌈M/32⌉·⌈N/32⌉).

pub fn gather_qmm( &self, w: &Quantized, lhs_indices: Option<&Array>, rhs_indices: Option<&Array>, options: GatherQmm<'_>, ) -> Array

Quantized products of gathered matrices (MLX’s gather_qmm). Up to ROW_INDEPENDENT_ROWS tokens of up to 8 experts run one product per row, sorted or not, each with the bits of its own product.

Memory

With sorted indices over many rows: where each expert’s rows start (an int32 per expert); otherwise none.

pub fn view(&self, dtype: DType) -> Array

The bits reinterpreted as dtype (MLX’s view): the same shape when the element sizes agree, else the last dimension scaled by their ratio (it must divide). Shares the buffer when the array is row-contiguous (or the sizes agree); copies otherwise.

use nuvec_core::{DType, array};

let bytes = array![[1u8, 0, 0, 0, 2, 0, 0, 0]];
assert_eq!(bytes.view(DType::U32).to_vec::<u32>().unwrap(), [1, 2]);

pub fn to_fp8(&self) -> Array

The floats as E4M3 bytes (uint8; MLX’s to_fp8): rounded to nearest even, saturating at ±448. There is no fp8 dtype; Array::from_fp8 restores floats.

use nuvec_core::{DType, array};

let codes = array![1.0, 448.0, 500.0, -3.3].to_fp8();
assert_eq!(codes.to_vec::<u8>().unwrap(), [56, 126, 126, 197]);
assert_eq!(codes.from_fp8(DType::F32).to_vec::<f32>().unwrap()[1], 448.0);

pub fn from_fp8(&self, dtype: DType) -> Array

E4M3 bytes (uint8) as dtype floats (MLX’s from_fp8).

pub fn gather_mm( &self, b: &Array, lhs_indices: Option<&Array>, rhs_indices: Option<&Array>, sorted_indices: bool, ) -> Array

Products of gathered matrices (MLX’s gather_mm). One row per output batch (up to ROW_INDEPENDENT_ROWS tokens of up to 8 experts) runs MLX’s gathered matrix-vector kernels, each row with the bits of its own product.

Memory

As Array::gather_qmm: the sorted path’s offsets; otherwise none (operands that are not row-major are copied first).

pub fn conv1d( &self, weight: &Array, stride: i32, padding: i32, dilation: i32, groups: i32, ) -> Array

A 1-D convolution of self ([N, L, C_in]) with weight ([C_out, K, C_in / groups]).

use nuvec_core::array;

let x = array![[[1.0], [2.0], [3.0]]]; // [1, 3, 1]
let w = array![[[1.0], [1.0]]]; // [1, 2, 1]
assert_eq!(x.conv1d(&w, 1, 0, 1, 1).to_vec::<f32>().unwrap(), [3.0, 5.0]);

pub fn zeros_like(&self) -> Array

Zeros of this array’s shape and dtype.

use nuvec_core::array;

assert_eq!(array![1u8, 2].zeros_like().to_vec::<u8>().unwrap(), [0, 0]);

pub fn ones_like(&self) -> Array

Ones of this array’s shape and dtype.

pub fn full_like(&self, value: impl Into<ScalarOrArray>) -> Array

value (a weak scalar takes this array’s dtype) in this array’s shape and dtype.

use nuvec_core::array;

assert_eq!(array![1.0, 2.0].full_like(7).to_vec::<f32>().unwrap(), [7.0, 7.0]);

pub fn t(&self) -> Array

The array with all axes reversed (the matrix transpose for 2-D arrays). Shares the data.

use nuvec_core::array;

assert_eq!(array![[1, 2, 3]].t().shape(), [3, 1]);

pub fn split_n<const N: usize>(&self, axis: impl Axis) -> [Array; N]

The array split into N equal parts along axis, for destructuring: let [q, gate] = qg.split_n::<2>(-1);. N failed arrays if the axis does not divide.

use nuvec_core::array;

let qg = array![[1.0, 2.0, 3.0, 4.0]];
let [q, g] = qg.split_n::<2>(-1);
assert_eq!((q.shape(), g.to_vec::<f32>().unwrap()), (&[1, 2][..], vec![3.0, 4.0]));

Trait Implementations

impl Add for Param

type Output = Array

The resulting type after applying the + operator.

fn add(self, rhs: Param) -> Array

Performs the + operation.

impl Add<&Array> for &Param

type Output = Array

The resulting type after applying the + operator.

fn add(self, rhs: &Array) -> Array

Performs the + operation.

impl Add<&Array> for Param

type Output = Array

The resulting type after applying the + operator.

fn add(self, rhs: &Array) -> Array

Performs the + operation.

impl Add<&Param> for &Array

type Output = Array

The resulting type after applying the + operator.

fn add(self, rhs: &Param) -> Array

Performs the + operation.

impl Add<&Param> for Array

type Output = Array

The resulting type after applying the + operator.

fn add(self, rhs: &Param) -> Array

Performs the + operation.

impl Add<&Param> for &Param

type Output = Array

The resulting type after applying the + operator.

fn add(self, rhs: &Param) -> Array

Performs the + operation.

impl Add<Array> for &Param

type Output = Array

The resulting type after applying the + operator.

fn add(self, rhs: Array) -> Array

Performs the + operation.

impl Add<Array> for Param

type Output = Array

The resulting type after applying the + operator.

fn add(self, rhs: Array) -> Array

Performs the + operation.

impl Add<Param> for &Array

type Output = Array

The resulting type after applying the + operator.

fn add(self, rhs: Param) -> Array

Performs the + operation.

impl Add<Param> for Array

type Output = Array

The resulting type after applying the + operator.

fn add(self, rhs: Param) -> Array

Performs the + operation.

impl<S: ScalarOperand> Add<S> for &Param

type Output = Array

The resulting type after applying the + operator.

fn add(self, rhs: S) -> Array

Performs the + operation.

impl<S: ScalarOperand> Add<S> for Param

type Output = Array

The resulting type after applying the + operator.

fn add(self, rhs: S) -> Array

Performs the + operation.

impl AsRef<Array> for Param

fn as_ref(&self) -> &Array

Converts this type into a shared reference of the (usually inferred) input type.

impl Clone for Param

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 Param

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl Default for Param

fn default() -> Self

Returns the “default value” for a type.

impl Deref for Param

type Target = Array

The resulting type after dereferencing.

fn deref(&self) -> &Array

Dereferences the value.

impl Display for Param

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl Div for Param

type Output = Array

The resulting type after applying the / operator.

fn div(self, rhs: Param) -> Array

Performs the / operation.

impl Div<&Array> for &Param

type Output = Array

The resulting type after applying the / operator.

fn div(self, rhs: &Array) -> Array

Performs the / operation.

impl Div<&Array> for Param

type Output = Array

The resulting type after applying the / operator.

fn div(self, rhs: &Array) -> Array

Performs the / operation.

impl Div<&Param> for &Array

type Output = Array

The resulting type after applying the / operator.

fn div(self, rhs: &Param) -> Array

Performs the / operation.

impl Div<&Param> for Array

type Output = Array

The resulting type after applying the / operator.

fn div(self, rhs: &Param) -> Array

Performs the / operation.

impl Div<&Param> for &Param

type Output = Array

The resulting type after applying the / operator.

fn div(self, rhs: &Param) -> Array

Performs the / operation.

impl Div<Array> for &Param

type Output = Array

The resulting type after applying the / operator.

fn div(self, rhs: Array) -> Array

Performs the / operation.

impl Div<Array> for Param

type Output = Array

The resulting type after applying the / operator.

fn div(self, rhs: Array) -> Array

Performs the / operation.

impl Div<Param> for &Array

type Output = Array

The resulting type after applying the / operator.

fn div(self, rhs: Param) -> Array

Performs the / operation.

impl Div<Param> for Array

type Output = Array

The resulting type after applying the / operator.

fn div(self, rhs: Param) -> Array

Performs the / operation.

impl<S: ScalarOperand> Div<S> for &Param

type Output = Array

The resulting type after applying the / operator.

fn div(self, rhs: S) -> Array

Performs the / operation.

impl<S: ScalarOperand> Div<S> for Param

type Output = Array

The resulting type after applying the / operator.

fn div(self, rhs: S) -> Array

Performs the / operation.

impl<'a> From<&'a Param> for Option<&'a Array>

A parameter where an optional array is expected (fast::rms_norm(x, &p, eps)).

fn from(p: &'a Param) -> Self

Converts to this type from the input type.

impl From<&Param> for ScalarOrArray

fn from(p: &Param) -> Self

Converts to this type from the input type.

impl From<Array> for Param

fn from(value: Array) -> Self

Converts to this type from the input type.

impl From<Param> for Array

fn from(p: Param) -> Self

Converts to this type from the input type.

impl From<Param> for ScalarOrArray

fn from(p: Param) -> Self

Converts to this type from the input type.

impl Leaves for Param

fn visit(&self, v: &mut Visit<'_>)

Reports every leaf (and child tree, enum variant, mode flag and static field) to v, in order.

impl Mul for Param

type Output = Array

The resulting type after applying the * operator.

fn mul(self, rhs: Param) -> Array

Performs the * operation.

impl Mul<&Array> for &Param

type Output = Array

The resulting type after applying the * operator.

fn mul(self, rhs: &Array) -> Array

Performs the * operation.

impl Mul<&Array> for Param

type Output = Array

The resulting type after applying the * operator.

fn mul(self, rhs: &Array) -> Array

Performs the * operation.

impl Mul<&Param> for &Array

type Output = Array

The resulting type after applying the * operator.

fn mul(self, rhs: &Param) -> Array

Performs the * operation.

impl Mul<&Param> for Array

type Output = Array

The resulting type after applying the * operator.

fn mul(self, rhs: &Param) -> Array

Performs the * operation.

impl Mul<&Param> for &Param

type Output = Array

The resulting type after applying the * operator.

fn mul(self, rhs: &Param) -> Array

Performs the * operation.

impl Mul<Array> for &Param

type Output = Array

The resulting type after applying the * operator.

fn mul(self, rhs: Array) -> Array

Performs the * operation.

impl Mul<Array> for Param

type Output = Array

The resulting type after applying the * operator.

fn mul(self, rhs: Array) -> Array

Performs the * operation.

impl Mul<Param> for &Array

type Output = Array

The resulting type after applying the * operator.

fn mul(self, rhs: Param) -> Array

Performs the * operation.

impl Mul<Param> for Array

type Output = Array

The resulting type after applying the * operator.

fn mul(self, rhs: Param) -> Array

Performs the * operation.

impl<S: ScalarOperand> Mul<S> for &Param

type Output = Array

The resulting type after applying the * operator.

fn mul(self, rhs: S) -> Array

Performs the * operation.

impl<S: ScalarOperand> Mul<S> for Param

type Output = Array

The resulting type after applying the * operator.

fn mul(self, rhs: S) -> Array

Performs the * operation.

impl Neg for &Param

type Output = Array

The resulting type after applying the - operator.

fn neg(self) -> Array

Performs the unary - operation.

impl Sub for Param

type Output = Array

The resulting type after applying the - operator.

fn sub(self, rhs: Param) -> Array

Performs the - operation.

impl Sub<&Array> for &Param

type Output = Array

The resulting type after applying the - operator.

fn sub(self, rhs: &Array) -> Array

Performs the - operation.

impl Sub<&Array> for Param

type Output = Array

The resulting type after applying the - operator.

fn sub(self, rhs: &Array) -> Array

Performs the - operation.

impl Sub<&Param> for &Array

type Output = Array

The resulting type after applying the - operator.

fn sub(self, rhs: &Param) -> Array

Performs the - operation.

impl Sub<&Param> for Array

type Output = Array

The resulting type after applying the - operator.

fn sub(self, rhs: &Param) -> Array

Performs the - operation.

impl Sub<&Param> for &Param

type Output = Array

The resulting type after applying the - operator.

fn sub(self, rhs: &Param) -> Array

Performs the - operation.

impl Sub<Array> for &Param

type Output = Array

The resulting type after applying the - operator.

fn sub(self, rhs: Array) -> Array

Performs the - operation.

impl Sub<Array> for Param

type Output = Array

The resulting type after applying the - operator.

fn sub(self, rhs: Array) -> Array

Performs the - operation.

impl Sub<Param> for &Array

type Output = Array

The resulting type after applying the - operator.

fn sub(self, rhs: Param) -> Array

Performs the - operation.

impl Sub<Param> for Array

type Output = Array

The resulting type after applying the - operator.

fn sub(self, rhs: Param) -> Array

Performs the - operation.

impl<S: ScalarOperand> Sub<S> for &Param

type Output = Array

The resulting type after applying the - operator.

fn sub(self, rhs: S) -> Array

Performs the - operation.

impl<S: ScalarOperand> Sub<S> for Param

type Output = Array

The resulting type after applying the - operator.

fn sub(self, rhs: S) -> Array

Performs the - operation.

impl Tree for Param

fn visit_mut(&mut self, v: &mut VisitMut<'_>)

Reports every leaf (and child tree and mode flag) to v mutably, in the order of Leaves::visit.

fn for_each(&self, f: impl FnMut(&str, &Array))

Calls f(path, leaf) for every leaf, in order.

fn map(&self, f: impl FnMut(&str, &Array) -> Array) -> Self

The tree with every leaf replaced by f(path, leaf) (same structure; parameters keep their frozen flag).

fn leaves(&self) -> Vec<Array>

The leaves, in order.

fn paths(&self) -> Vec<String>

The leaves’ paths, in order.

fn flatten(&self) -> Vec<(String, Array)>

The leaves with their paths, in order (MLX’s tree_flatten).

fn unflatten(&self, leaves: impl IntoIterator<Item = Array>) -> Result<Self>

A tree of this structure (self is the template) with leaves in order: exactly Tree::len of them, else an error naming both counts.

fn load_flat( &mut self, named: impl IntoIterator<Item = (String, Array)>, mode: LoadMode, ) -> Result<()>

Replaces leaves by path. LoadMode::Strict requires exactly the tree’s paths with matching shapes; LoadMode::Partial assigns the matching names and ignores unknown ones (shapes must still match). Loaded dtypes are kept. On any problem the error lists every one (missing, unknown and mismatched names, sorted) and the tree is unchanged. Duplicate paths (in the tree or in named) are an error.

fn len(&self) -> usize

The number of leaves.

fn is_empty(&self) -> bool

Whether the tree has no leaves.

fn def(&self) -> TreeDef

The tree’s structure.