rlx-tensor 0.2.8

Native ndarray alternative for RLX — NumPy-style tensors that trace, fuse, and run on any backend via the rlx compiler
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
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155

// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.

//! # rlx-tensor
//!
//! Symbolic tensor API for RLX — expression syntax with compiler-backed
//! execution. Compared to NumPy / ndarray / PyTorch / JAX:
//!
//! | | ndarray / NumPy | PyTorch | JAX | **rlx-tensor** |
//! |-|-----------------|---------|-----|----------------|
//! | Default execution | eager, host | eager, GPU | trace + XLA | **trace → fuse → backend** |
//! | Intermediate buffers | always materialized | often kept for autograd | XLA decides | **fusion + memory plan** |
//! | Device | host only | CUDA-centric | TPU/GPU | **CPU / Metal / CUDA / ROCm / wgpu / MLX / TPU** |
//! | Variable batch/seq | manual padding | manual | `jax.vmap` / shapes | **first-class `Dim::Dynamic`** |
//! | Slicing | views (cheap) | views | slice ops | **`s![]` → narrow (fusable)** |
//!
//! Symbolic [`Tensor`] handles carry **zero payload bytes** until you compile
//! and run — but with the `eval` feature you can also build from host data
//! NumPy-style and read results straight back, the graph staying lazy (so the
//! compiler fuses + memory-plans the whole expression) until you materialize:
//!
//! ```ignore
//! use rlx_tensor::Tensor; // requires feature = "eval"
//!
//! let a = Tensor::from_vec(vec![1.0, 2.0, 3.0], [3]);
//! let b = Tensor::ones([3]);
//! let c = (&a + &b).relu();        // graphs auto-merge, nothing runs yet
//! assert_eq!(c.to_vec(), vec![2.0, 3.0, 4.0]); // auto-selects fastest backend
//! // c.on(Device::Metal).to_vec()  // or pin an explicit backend
//! ```
//!
//! Reverse-mode autodiff (feature `grad`) returns gradient tensors:
//!
//! ```ignore
//! let loss = (&a * &b).sum([0], false);
//! let g = loss.grad(&[&a, &b]);     // [∂loss/∂a, ∂loss/∂b]
//! assert_eq!(g[0].to_vec(), b.to_vec());
//! ```
//!
//! JAX-shaped composable transforms (feature `transforms`) operate on [`Func`]
//! (a traced function of named inputs) and chain — `vmap(grad(f))`:
//!
//! ```ignore
//! let f = Func::new("f", |s| { let x = s.input("x", shape![3]); (&x * &x).sum([0], false) });
//! let batched_grad = f.grad(&["x"]).vmap(&["x"], 4);  // Func -> Func -> Func
//! let compiled = batched_grad.jit();                  // compile once
//! let out = compiled.run(&[("x", &xs)]);              // run many, no recompile
//! ```
//!
//! Materialization is compile-cached per thread (keyed by graph + constants +
//! device), so even plain `to_vec` / `run` reuse the compiled artifact.
//!
//! ```rust
//! use rlx_tensor::{ax, graph, rg, s, shape};
//!
//! // static shape
//! let g = graph("mlp", |g| {
//!     let x = g.input("x", shape![2, 4]);
//!     let w = g.param("w", shape![4, 3]);
//!     (&x.matmul(&w)).gelu()
//! });
//!
//! // dynamic batch + window slice (no copy in IR)
//! let g2 = graph("batched", |g| {
//!     let x = g.input("x", shape![?, 128]);
//!     x.slice(s![ax(), rg(0, 64)])
//! });
//! # let _ = (g, g2);
//! ```

mod array;
#[cfg(feature = "eval")]
mod cache;
mod handle;
#[cfg(feature = "ndarray")]
mod interop;
mod scalar;
#[cfg(feature = "optim")]
mod schedule;
mod scope;
mod slice;
mod tensor;
mod transform;

pub use array::{cat, stack};
pub use rlx_ir::op::MaskKind;
pub use rlx_ir::{DType, Dim, Graph, NodeId, Shape};
/// Optimizers for [`Func::train_step`] (re-exported from `rlx_optim`).
/// Available with the `optim` feature.
#[cfg(feature = "optim")]
pub use rlx_optim::{Adam, AdamW, Lion, Optimizer, Sgd};
#[cfg(feature = "optim")]
pub use schedule::LrSchedule;
pub use scope::{GraphScope, graph, graph_with};
pub use slice::{SliceAxis, SliceSpec, ax, ix, rg, tail};
pub use tensor::Tensor;
pub use transform::Func;
#[cfg(feature = "eval")]
pub use transform::Jitted;

#[cfg(feature = "eval")]
pub use cache::{cache_stats, clear_cache};
/// Device selector for [`Tensor::on`] / [`Tensor::to_vec_on`] (re-exported
/// from `rlx_runtime`). Available with the `eval` feature.
#[cfg(feature = "eval")]
pub use rlx_runtime::Device;
/// Backend detection (re-exported from `rlx_runtime`): which devices this build
/// can use, and the fastest one. `to_vec` / `Func::run` pick automatically.
#[cfg(feature = "eval")]
pub use rlx_runtime::{available_devices, fastest_device, fastest_device_for, is_available};
#[cfg(feature = "eval")]
pub use tensor::Materialize;

#[doc(hidden)]
#[macro_export]
macro_rules! __dim {
    (?) => {
        $crate::Dim::Dynamic(0)
    };
    (? $sym:literal) => {
        $crate::Dim::Dynamic($sym)
    };
    ($n:expr) => {
        $crate::Dim::Static($n)
    };
}

/// Build a [`Shape`]. Default dtype is `F32`.
///
/// Dynamic dimensions: `?` → symbol 0 (repeat `?` for the same unknown size),
/// `?1` → symbol 1, etc.
///
/// ```rust
/// use rlx_tensor::{shape, DType};
///
/// let a = shape![2, 4];
/// let b = shape![F32; ?, 128];
/// assert_eq!(b.dtype(), DType::F32);
/// assert!(!b.is_static());
/// ```
#[macro_export]
macro_rules! shape {
    ($dtype:ident; $($d:tt),+ $(,)?) => {
        $crate::Shape::from_dims(
            &[$( $crate::__dim!($d) ),+],
            $crate::DType::$dtype,
        )
    };
    ($($d:tt),+ $(,)?) => {
        $crate::Shape::from_dims(
            &[$( $crate::__dim!($d) ),+],
            $crate::DType::F32,
        )
    };
}