Crate rlx_tensor

Expand description

§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:

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:

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)):

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.

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)])
});

Or, with the dsl feature, declare the whole graph in a compact little language with the rlx! macro — @ is matmul, shapes and outputs are inferred.

Modules§

lowp: Generic low-precision numerics — the float_format! macro + fXmYeZ quantizer for emulated any-precision (nvf4 / f8 / bf8 / …) training. Generic low-precision numerics for emulated mixed-precision training — quantize an f32 to the grid of an arbitrary binary float format fXmYeZ (X total bits, Y mantissa bits, Z exponent bits, 1 sign), then keep computing in f32. This makes training work at any precision on any backend (including exotic formats with no hardware kernel — nvf4, f8, bf8) by round-tripping values through the format’s representable set instead of relying on native half/quarter-precision matmul.

Macros§

float_format: Generate a marker type + FloatFormat impl per format, plus a BUILTIN lookup table. $name = (exp_bits, man_bits, max_normal, "spec").
rlx: Declare a computation Graph in a compact, readable little language.
rlx_expr: Evaluate a single rlx!-grammar expression over in-scope Rust Tensor values — the “Rust bridge”. Every identifier is a normal Rust variable, so ordinary for loops and config drive the graph structure while the DSL expresses the per-step math (e.g. h = rlx_expr!(gelu(linear(h, w, b)))). This is what a macro-expanded rlx! { … } can’t do: runtime-length loops. Requires the dsl feature.
s: Build a SliceSpec from axis expressions.
shape: Build a Shape. Default dtype is F32.

Structs§

Adam: Optimizers for Func::train_step (re-exported from rlx_optim). Available with the optim feature. Bias-corrected first/second moment optimizer.
AdamW: Optimizers for Func::train_step (re-exported from rlx_optim). Available with the optim feature. Adam with decoupled weight decay.
Func: A function of named inputs and parameters, traced into a graph. Unit of composition for the grad/vmap/ jvp/hvp transforms.
Graph: A computation graph — the core IR data structure.
GraphScope: Owns a traced graph while you build symbolic tensors.
Jitted: An already-compiled Func, returned by Func::jit. Cloning is cheap (shared compiled graph); calls run on the same backend artifact.
Lion: Optimizers for Func::train_step (re-exported from rlx_optim). Available with the optim feature. EvoLved sign-momentum optimizer.
Materialize: Fluent device selector returned by Tensor::on.
Model: A compiled, runnable model with by-name weight binding. Available with the eval feature. A compiled graph plus its parameter roster — a runnable model.
Muon: Optimizers for Func::train_step (re-exported from rlx_optim). Available with the optim feature. Muon — Momentum-Orthogonalized-by-Newton-Schulz.
NodeId: Stable identifier for a node in the graph. Indices are never reused.
OptItem: Optimizers for Func::train_step (re-exported from rlx_optim). Available with the optim feature. Common parameter-update interface.
Sgd: Optimizers for Func::train_step (re-exported from rlx_optim). Available with the optim feature. SGD with momentum / Nesterov / L2 weight decay.
Shape: Tensor shape: ordered list of dimensions + element type.
SliceSpec: Multi-axis slice. Apply via Tensor::slice.
Tensor: A symbolic tensor in a traced graph.

Enums§

Activation: Unary element-wise activation functions.
DType: Scalar element type. Matches hardware-supported types.
Device: Device selector for Tensor::on / Tensor::to_vec_on (re-exported from rlx_runtime). Available with the eval feature. Target device for graph execution.
Dim: A single dimension — either a concrete size or a symbolic dynamic dim.
LrSchedule: Common learning-rate schedules.
MaskKind: What kind of attention mask the kernel should apply.
ScaleLayout: How scale factors are laid out for an crate::op::Op::ScaledMatMul operand. The reconstructed value of element i is decode(code[i]) * scale(block_of(i)).
ScaledFormat: Element format for a native low-precision tensor-core GEMM operand (see crate::op::Op::ScaledMatMul).
SliceAxis: One axis in a slice (s![ax(), rg(2, 10)] or s![.., 2..10] via helpers).

Traits§

Optimizer: Optimizers for Func::train_step (re-exported from rlx_optim). Available with the optim feature.

Functions§

available_devices: Backend detection (re-exported from rlx_runtime): which devices this build can use, and the fastest one. to_vec / Func::run pick automatically. Every variant currently available — Cargo-feature-gated or runtime-registered.
ax: Keep the full axis.
cache_stats: (hits, misses) for the calling thread’s cache. Useful for tests and confirming jit-style reuse.
cat: Concatenate tensors along an existing axis (NumPy concatenate). Inputs from different graphs are merged transparently.
clear_cache: Drop all cached compiled graphs and reset stats on the calling thread.
fastest_device: Backend detection (re-exported from rlx_runtime): which devices this build can use, and the fastest one. to_vec / Func::run pick automatically. Highest-priority backend that is compiled in and live on this host.
fastest_device_for: Backend detection (re-exported from rlx_runtime): which devices this build can use, and the fastest one. to_vec / Func::run pick automatically. Pick the fastest backend for graph on this host.
graph: Build a single-output graph from a closure.
graph_with: Build a graph and return extra metadata from the closure.
is_available: Backend detection (re-exported from rlx_runtime): which devices this build can use, and the fastest one. to_vec / Func::run pick automatically.
ix: Single index (negative counts from the end). Keeps the axis as size 1 — use Tensor::select to drop it.
rg: Half-open range start..end (either bound may be negative).
stack: Stack tensors along a new axis (NumPy stack): each gains a size-1 axis at axis, then they are concatenated there.
tail: start.. through the end of a static axis (negative start allowed).

Crate rlx_tensor

Crate rlx_tensor Copy item path

§rlx-tensor

Modules§

Macros§

Structs§

Enums§

Traits§

Functions§

Crate rlx_tensor