rlx

A small ML compiler and runtime for transformer inference and training. JAX-shaped IR + autodiff + transforms (jvp, hvp, vmap) on top of backend-specific kernels for CPU, Apple Silicon (Metal / MLX), NVIDIA (CUDA), AMD (ROCm), Google TPU, and cross-platform GPU (wgpu).

This is the prelude crate — pulls in rlx-ir / rlx-opt / rlx-runtime and re-exports the common types. Most code only needs one use rlx::prelude::*;.

Install

[dependencies]
rlx = { version = "0.2", features = ["cpu"] }

For common platforms, single-flag aggregates compose the right fragments:

rlx = { version = "0.2", features = ["apple-silicon"] }   # cpu + metal + Accelerate
rlx = { version = "0.2", features = ["nvidia"] }          # cpu + cuda
rlx = { version = "0.2", features = ["edge"] }            # cpu + cortexm
rlx = { version = "0.2", features = ["all-cpu"] }         # cpu + gguf + linalg

mlx and rocm features. rlx-mlx and rlx-rocm aren't on crates.io (vendor-bundled submodule / workspace-relative kernel sources). Enabling those features on a crates.io build will fail to resolve. Use a git source instead:

rlx = { git = "https://github.com/MIT-RLX/rlx", features = ["apple-silicon", "mlx"] }

Quickstart

use rlx::prelude::*;

let mut g = Graph::new("hello");
let x = g.input("x", Shape::new(&[1, 4], DType::F32));
let w = g.param("w", Shape::new(&[4, 2], DType::F32));
let y = g.matmul(x, w, Shape::new(&[1, 2], DType::F32));
g.set_outputs(vec![y]);

let mut compiled = Session::new(Device::Cpu).compile(g);
compiled.set_param("w", &[1.0, 0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 1.0]);
let out = compiled.run(&[("x", &[1.0, 2.0, 3.0, 4.0])]);

Prefer to write the math? The rlx! DSL (feature tensor, on by default) builds the same graph in a compact little language — @ is matmul, shapes and outputs are inferred:

use rlx::rlx;

let g = rlx! {
    input x: [1, 4];
    param w: [4, 2];
    let y = x @ w;      // → rlx_ir::Graph
};

Multi-backend runtime

let mut runner = GraphDevices::with_policy(g, DevicePolicy::only([Device::Cpu, Device::Metal]));
let out = runner.run_resolved_with_inputs(None, &inputs)?;

let mut router = DeviceRouter::from_env(g)?;
let (device, out) = router.run(&inputs, None)?;

See docs/backend-selection.md.

Prelude + namespaces

With features = ["optim"], the prelude also star-imports Adam, AdamW, Optimizer, and the rest of rlx-optim.

rlx::Result<T> and rlx::Error are aliases of anyhow::Result<T> and anyhow::Error — the whole stack returns those.

Feature matrix

Backends

Companion crates

Off by default; turn on per workload:

cortexm kernels and fpga export don't go through Session / Backend execution — they're specialty / offline targets. FPGA RTL is target-agnostic by default (HwTarget::Generic); see docs/fpga-export.md.

Convenience aggregates

mlx and rocm aren't in any aggregate (vendor-bundled). To opt in, add the feature explicitly to a git-source dep:

rlx = { git = "https://github.com/MIT-RLX/rlx", features = ["apple-silicon", "mlx"] }

Documentation

• API reference: https://docs.rs/rlx
• Workspace overview + per-crate READMEs: https://github.com/MIT-RLX/rlx

License

MIT OR Apache-2.0, at your option. See LICENSE-MIT and LICENSE-APACHE.