Crate rlx_tensor 

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

Macros

s

Build a SliceSpec from axis expressions.

shape

Build a Shape. Default dtype is F32.

Structs

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.

NodeId

Stable identifier for a node in the graph. Indices are never reused.

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

DType

Scalar element type. Matches hardware-supported types.

Dim

A single dimension — either a concrete size or a symbolic dynamic dim.

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

Functions

ax

Keep the full axis.

cat

Concatenate tensors along an existing axis (NumPy concatenate). Inputs from different graphs are merged transparently.

graph

Build a single-output graph from a closure.

graph_with

Build a graph and return extra metadata from the closure.

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