nuvec_core/
lib.rs

1//! Lazy array graph, primitives, dtypes, streams and evaluation: the core of nuvec.
2//!
3//! Operations on [`Array`]s record a graph and return immediately. Nothing is computed
4//! until an array is evaluated, explicitly or by reading it back:
5//!
6//! ```
7//! use nuvec_core::{Array, Result};
8//!
9//! # fn main() -> Result<()> {
10//! let a = Array::from_slice(&[1.0f32, 2.0, 3.0], &[3])?;
11//! let b = Array::full(&[3], 10.0f32)?;
12//! let c = (&a + &b) * &a; // lazy: nothing computed yet
13//! assert!(!c.is_evaluated());
14//! assert_eq!(c.to_vec::<f32>()?, vec![11.0, 24.0, 39.0]);
15//! # Ok(())
16//! # }
17//! ```
18//!
19//! Binary operations promote types (MLX's table) and broadcast shapes (NumPy rules); views
20//! (`reshape`, `transpose`, `slice`, `broadcast_to`, `expand_dims`, `squeeze`) share data.
21//!
22//! Reductions (`sum`, `mean`, `max`, `argmax`, …) work over any axes and layouts;
23//! `matmul` broadcasts batch dimensions and uses a SIMD GEMM on the CPU.
24//!
25//! Arrays live in unified memory. Operations on the GPU stream run as Metal kernels; anything
26//! without a GPU kernel (for example float64) runs on the CPU over the same memory:
27//!
28//! ```
29//! use nuvec_core::{Array, Device, Stream, set_default_stream};
30//!
31//! # fn main() -> nuvec_core::Result<()> {
32//! set_default_stream(Stream::default_for(Device::gpu(0)));
33//! let x = Array::from_slice(&[1.0f32, 2.0, 3.0], &[3])?;
34//! let y = &x * &x; // a Metal kernel on a Mac; an `UnsupportedDevice` error elsewhere
35//! # if y.eval().is_ok() { assert_eq!(y.to_vec::<f32>()?, vec![1.0, 4.0, 9.0]); }
36//! # set_default_stream(Stream::default_for(Device::cpu()));
37//! # Ok(())
38//! # }
39//! ```
40//!
41//! **Status:** pre-alpha. The default stream is the CPU.
42
43#![forbid(unsafe_code)]
44
45mod array;
46mod batching;
47mod buffer;
48mod cast;
49mod conv;
50mod custom;
51mod derivatives;
52mod device;
53mod dtype;
54mod error;
55mod eval;
56pub mod fast;
57mod file;
58// Parts of it are only used by the Metal backend.
59#[cfg_attr(not(all(target_os = "macos", not(miri))), allow(dead_code))]
60mod fused;
61mod gpu;
62mod indexing;
63// Parts of it are only used by the Metal backend.
64#[cfg_attr(not(all(target_os = "macos", not(miri))), allow(dead_code))]
65mod kernel;
66mod kernels;
67mod ops;
68mod primitive;
69mod quantized;
70pub mod random;
71mod scalar;
72mod shape;
73mod sorting;
74mod transforms;
75mod view;
76
77pub use array::Array;
78pub use conv::conv1d;
79pub use custom::{CustomCall, CustomFunction, checkpoint, custom_function};
80pub use device::{Device, DeviceKind, Stream, default_device, default_stream, set_default_stream};
81pub use dtype::{Category, Complex64, DType, Element, at_least_float, bf16, f16, promote_types};
82pub use error::{Error, Result};
83pub use eval::{Scheduled, async_eval, eval, schedule};
84pub use file::MappedFile;
85pub use fused::{Fused, FusedRef};
86pub use indexing::{
87    concatenate, put_along_axis, scatter_add_along_axis, split, split_at, stack, take,
88    take_along_axis,
89};
90pub use ops::{
91    abs, add, addmm, all, any, arange, argmax, argmin, astype, broadcast_to, conjugate, contiguous,
92    cos, divide, equal, erfinv, exp, expand_dims, greater, greater_equal, less, less_equal, log,
93    logaddexp, logsumexp, matmul, max, maximum, mean, min, minimum, multiply, negative, not_equal,
94    ones, prod, repeat, reshape, sigmoid, sign, sin, slice, slice_update, softmax, sqrt, squeeze,
95    stop_gradient, subtract, sum, tanh, transpose, where_, zeros,
96};
97pub use primitive::{ArgReduceOp, BinaryOp, Primitive, ReduceOp, ScanOp, ScatterReduce, UnaryOp};
98pub use quantized::{dequantize, gather_mm, gather_qmm, quantize, quantized_matmul};
99pub use scalar::{Scalar, ScalarOperand, ScalarOrArray};
100pub use shape::Flags;
101pub use sorting::{
102    argpartition, argsort, cummax, cummin, cumprod, cumsum, logcumsumexp, partition, sort, topk,
103};
104pub use transforms::{grad, jvp, value_and_grad, vjp, vmap};
105pub use view::SliceArg;
106
107/// The version of this crate.
108pub const VERSION: &str = env!("CARGO_PKG_VERSION");
nuvec_core/lib.rs

nuvec_core/
lib.rs
