rlx_tensor/
lib.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3
4//! # rlx-tensor
5//!
6//! Symbolic tensor API for RLX — expression syntax with compiler-backed
7//! execution. Compared to NumPy / ndarray / PyTorch / JAX:
8//!
9//! | | ndarray / NumPy | PyTorch | JAX | **rlx-tensor** |
10//! |-|-----------------|---------|-----|----------------|
11//! | Default execution | eager, host | eager, GPU | trace + XLA | **trace → fuse → backend** |
12//! | Intermediate buffers | always materialized | often kept for autograd | XLA decides | **fusion + memory plan** |
13//! | Device | host only | CUDA-centric | TPU/GPU | **CPU / Metal / CUDA / ROCm / wgpu / MLX / TPU** |
14//! | Variable batch/seq | manual padding | manual | `jax.vmap` / shapes | **first-class `Dim::Dynamic`** |
15//! | Slicing | views (cheap) | views | slice ops | **`s![]` → narrow (fusable)** |
16//!
17//! Symbolic [`Tensor`] handles carry **zero payload bytes** until you compile
18//! and run — but with the `eval` feature you can also build from host data
19//! NumPy-style and read results straight back, the graph staying lazy (so the
20//! compiler fuses + memory-plans the whole expression) until you materialize:
21//!
22//! ```ignore
23//! use rlx_tensor::Tensor; // requires feature = "eval"
24//!
25//! let a = Tensor::from_vec(vec![1.0, 2.0, 3.0], [3]);
26//! let b = Tensor::ones([3]);
27//! let c = (&a + &b).relu();        // graphs auto-merge, nothing runs yet
28//! assert_eq!(c.to_vec(), vec![2.0, 3.0, 4.0]); // auto-selects fastest backend
29//! // c.on(Device::Metal).to_vec()  // or pin an explicit backend
30//! ```
31//!
32//! Reverse-mode autodiff (feature `grad`) returns gradient tensors:
33//!
34//! ```ignore
35//! let loss = (&a * &b).sum([0], false);
36//! let g = loss.grad(&[&a, &b]);     // [∂loss/∂a, ∂loss/∂b]
37//! assert_eq!(g[0].to_vec(), b.to_vec());
38//! ```
39//!
40//! JAX-shaped composable transforms (feature `transforms`) operate on [`Func`]
41//! (a traced function of named inputs) and chain — `vmap(grad(f))`:
42//!
43//! ```ignore
44//! let f = Func::new("f", |s| { let x = s.input("x", shape![3]); (&x * &x).sum([0], false) });
45//! let batched_grad = f.grad(&["x"]).vmap(&["x"], 4);  // Func -> Func -> Func
46//! let compiled = batched_grad.jit();                  // compile once
47//! let out = compiled.run(&[("x", &xs)]);              // run many, no recompile
48//! ```
49//!
50//! Materialization is compile-cached per thread (keyed by graph + constants +
51//! device), so even plain `to_vec` / `run` reuse the compiled artifact.
52//!
53//! ```rust
54//! use rlx_tensor::{ax, graph, rg, s, shape};
55//!
56//! // static shape
57//! let g = graph("mlp", |g| {
58//!     let x = g.input("x", shape![2, 4]);
59//!     let w = g.param("w", shape![4, 3]);
60//!     (&x.matmul(&w)).gelu()
61//! });
62//!
63//! // dynamic batch + window slice (no copy in IR)
64//! let g2 = graph("batched", |g| {
65//!     let x = g.input("x", shape![?, 128]);
66//!     x.slice(s![ax(), rg(0, 64)])
67//! });
68//! # let _ = (g, g2);
69//! ```
70
71mod array;
72#[cfg(feature = "eval")]
73mod cache;
74mod handle;
75#[cfg(feature = "ndarray")]
76mod interop;
77mod scalar;
78#[cfg(feature = "optim")]
79mod schedule;
80mod scope;
81mod slice;
82mod tensor;
83mod transform;
84
85pub use array::{cat, stack};
86pub use rlx_ir::op::MaskKind;
87pub use rlx_ir::{DType, Dim, Graph, NodeId, ScaleLayout, ScaledFormat, Shape};
88/// Optimizers for [`Func::train_step`] (re-exported from `rlx_optim`).
89/// Available with the `optim` feature.
90#[cfg(feature = "optim")]
91pub use rlx_optim::{Adam, AdamW, Lion, Muon, Optimizer, Sgd};
92#[cfg(feature = "optim")]
93pub use schedule::LrSchedule;
94pub use scope::{GraphScope, graph, graph_with};
95pub use slice::{SliceAxis, SliceSpec, ax, ix, rg, tail};
96pub use tensor::Tensor;
97pub use transform::Func;
98#[cfg(feature = "eval")]
99pub use transform::Jitted;
100
101#[cfg(feature = "eval")]
102pub use cache::{cache_stats, clear_cache};
103/// Device selector for [`Tensor::on`] / [`Tensor::to_vec_on`] (re-exported
104/// from `rlx_runtime`). Available with the `eval` feature.
105#[cfg(feature = "eval")]
106pub use rlx_runtime::Device;
107/// Backend detection (re-exported from `rlx_runtime`): which devices this build
108/// can use, and the fastest one. `to_vec` / `Func::run` pick automatically.
109#[cfg(feature = "eval")]
110pub use rlx_runtime::{available_devices, fastest_device, fastest_device_for, is_available};
111#[cfg(feature = "eval")]
112pub use tensor::Materialize;
113
114#[doc(hidden)]
115#[macro_export]
116macro_rules! __dim {
117    (?) => {
118        $crate::Dim::Dynamic(0)
119    };
120    (? $sym:literal) => {
121        $crate::Dim::Dynamic($sym)
122    };
123    ($n:expr) => {
124        $crate::Dim::Static($n)
125    };
126}
127
128/// Build a [`Shape`]. Default dtype is `F32`.
129///
130/// Each static dimension may be **any `usize` expression** — a literal, an
131/// identifier, or arithmetic like `w / 2`, `w + heads`, `cfg.width * 2` — so graph
132/// builders can size params/inputs from computed dimensions directly (no need to
133/// hoist every derived size into a `let`). Dynamic dimensions use `?` → symbol 0
134/// (repeat `?` for the same unknown size), `?1` → symbol 1, etc. A leading
135/// `Dtype;` sets the element type (default `F32`).
136///
137/// ```rust
138/// use rlx_tensor::{shape, DType};
139///
140/// let a = shape![2, 4];
141/// let b = shape![F32; ?, 128];
142/// let w = 64usize;
143/// let c = shape![w + 1, w / 2, 3 * w];   // computed dimensions
144/// assert_eq!(c.dtype(), DType::F32);
145/// assert!(c.is_static());
146/// assert!(!b.is_static());
147/// ```
148#[macro_export]
149macro_rules! shape {
150    ($dtype:ident; $($rest:tt)+) => {
151        $crate::Shape::from_dims(&$crate::__shape_dims!([] ; $($rest)+), $crate::DType::$dtype)
152    };
153    ($($rest:tt)+) => {
154        $crate::Shape::from_dims(&$crate::__shape_dims!([] ; $($rest)+), $crate::DType::F32)
155    };
156}
157
158/// Token-tree muncher backing [`shape!`]: accumulates `Dim`s left to right,
159/// dispatching `?`/`?N` to dynamic dims and everything else to a `usize` expression.
160/// (A `tt`-list can't mix the `?` syntax with multi-token expressions, so we munch.)
161#[doc(hidden)]
162#[macro_export]
163macro_rules! __shape_dims {
164    // done (with or without a trailing comma)
165    ([$($acc:tt)*] ; ) => { [ $($acc)* ] };
166    ([$($acc:tt)*] ; ,) => { [ $($acc)* ] };
167    // dynamic `?N`
168    ([$($acc:tt)*] ; ? $sym:literal , $($rest:tt)*) => {
169        $crate::__shape_dims!([$($acc)* $crate::Dim::Dynamic($sym),] ; $($rest)*)
170    };
171    ([$($acc:tt)*] ; ? $sym:literal) => {
172        [ $($acc)* $crate::Dim::Dynamic($sym) ]
173    };
174    // dynamic `?` (symbol 0)
175    ([$($acc:tt)*] ; ? , $($rest:tt)*) => {
176        $crate::__shape_dims!([$($acc)* $crate::Dim::Dynamic(0),] ; $($rest)*)
177    };
178    ([$($acc:tt)*] ; ?) => {
179        [ $($acc)* $crate::Dim::Dynamic(0) ]
180    };
181    // static dimension: any expression, more following
182    ([$($acc:tt)*] ; $e:expr, $($rest:tt)*) => {
183        $crate::__shape_dims!([$($acc)* $crate::Dim::Static($e),] ; $($rest)*)
184    };
185    // static dimension: any expression, last
186    ([$($acc:tt)*] ; $e:expr) => {
187        [ $($acc)* $crate::Dim::Static($e) ]
188    };
189}
rlx_tensor/lib.rs

rlx_tensor/
lib.rs
