nuvec 0.1.0

Pure-Rust array framework for Apple Silicon in the spirit of MLX: lazy arrays, unified memory, composable transforms
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95

//! A pure-Rust array framework for Apple Silicon in the spirit of
//! [MLX](https://github.com/ml-explore/mlx): lazy arrays on unified memory and
//! composable function transforms, with no bindings to any C/C++ library.
//!
//! ```
//! use nuvec::{Array, DType, Result, arange};
//!
//! # fn main() -> Result<()> {
//! let x = arange(0.0, 6.0, 1.0, DType::F32)?.reshape(&[2, 3])?; // [[0,1,2],[3,4,5]]
//! let col = Array::from_slice(&[10i32, 20], &[2, 1])?;          // 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>()?, vec![10.0, 26.0, 12.0, 28.0, 14.0, 30.0]);
//! # Ok(())
//! # }
//! ```
//!
//! Gradients are lazy arrays too, so they compose:
//!
//! ```
//! use nuvec::{Array, Result, grad, multiply, sum};
//!
//! # fn main() -> Result<()> {
//! let f = |x: &Array| sum(&multiply(&multiply(x, x)?, x)?, &[], false); // sum(x³)
//! let x = Array::from_slice(&[1.0f32, 2.0], &[2])?;
//! assert_eq!(grad(&f)(&x)?.to_vec::<f32>()?, vec![3.0, 12.0]);          // 3x²
//! assert_eq!(grad(|x: &Array| sum(&grad(&f)(x)?, &[], false))(&x)?.to_vec::<f32>()?, vec![6.0, 12.0]);
//! # Ok(())
//! # }
//! ```
//!
//! 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`] and [`compile()`] (kernel fusion).
//! - **Metal kernels** ported from MLX (GEMM, GEMV, quantized products, attention,
//!   normalisation, RoPE, sort, scan), MLX-compatible affine [quantization](quantize) and
//!   custom kernels from source ([`fast::MetalKernel`]).
//! - **Evaluation** on the GPU by default, submitted asynchronously ([`async_eval`],
//!   [`schedule`]), or on the CPU; memory limits and wired memory on [`gpu::Device`].
//! - **I/O** ([`io`]): safetensors and npy, memory-mapped without a copy.
//!
//! The crate root re-exports the everyday API; the layers below are available as modules.
//! Neural-network layers and optimizers are in `nuvec-nn`, the LLM runtime in `nuvec-lm`.
//!
//! **Status:** pre-alpha: the API may change between releases.

#![forbid(unsafe_code)]

pub use nuvec_compile as compile;
pub use nuvec_compile::compile;
pub use nuvec_core as core;
pub use nuvec_core::{
    ArgReduceOp, Array, BinaryOp, Category, Complex64, CustomCall, CustomFunction, DType, Device,
    DeviceKind, Element, Error, Flags, Primitive, ReduceOp, Result, Scalar, ScalarOperand,
    ScalarOrArray, ScanOp, ScatterReduce, Scheduled, SliceArg, Stream, UnaryOp, abs, add, addmm,
    all, any, arange, argmax, argmin, argpartition, argsort, astype, async_eval, at_least_float,
    bf16, broadcast_to, checkpoint, concatenate, conjugate, contiguous, conv1d, cos, cummax,
    cummin, cumprod, cumsum, custom_function, default_device, default_stream, dequantize, divide,
    equal, erfinv, eval, exp, expand_dims, f16, gather_mm, gather_qmm, grad, greater,
    greater_equal, jvp, less, less_equal, log, logaddexp, logcumsumexp, logsumexp, matmul, max,
    maximum, mean, min, minimum, multiply, negative, not_equal, ones, partition, prod,
    promote_types, put_along_axis, quantize, quantized_matmul, repeat, reshape,
    scatter_add_along_axis, schedule, set_default_stream, sigmoid, sign, sin, slice, slice_update,
    softmax, sort, split, split_at, sqrt, squeeze, stack, stop_gradient, subtract, sum, take,
    take_along_axis, tanh, topk, transpose, value_and_grad, vjp, vmap, where_, zeros,
};
pub use nuvec_core::{fast, random};
pub use nuvec_cpu as cpu;
pub use nuvec_distributed as distributed;
pub use nuvec_gpu as gpu;
pub use nuvec_io as io;

/// The version of this crate.
pub const VERSION: &str = env!("CARGO_PKG_VERSION");

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn version_is_set() {
        assert!(!VERSION.is_empty());
        assert_eq!(VERSION, crate::core::VERSION);
    }

    #[test]
    fn root_api_works() {
        let a = Array::from_slice(&[1i32, 2], &[2]).unwrap();
        let b = multiply(&a, &a).unwrap() + &a;
        assert_eq!(b.to_vec::<i32>().unwrap(), vec![2, 6]);
    }
}