tracel-ai/burn

Burn – A Rust‑first tensor library and deep‑learning framework

What it is – Burn is an open‑source Rust crate that provides both low‑level tensor operations and a high‑level deep‑learning API. It aims to let you write the same Rust code for training a model and for running that model in production, removing the usual Python‑to‑ONNX export step.

Why it matters – Most research frameworks (PyTorch, TensorFlow) are Python‑centric; the trained model is later exported to a separate runtime (ONNX, TensorRT, etc.). Burn unifies these stages:

• Dynamic‑graph ergonomics similar to PyTorch (shapes can change at runtime).
• JIT‑compilation and automatic kernel fusion give GPU‑level speed while keeping the flexibility of a dynamic graph.
• Incremental compilation in Rust keeps edit‑run cycles short (re‑compile < 5 s even in release mode).

Key components

Area	Crate / Project	What it provides
Compute	CubeCL	A Rust‑based GPU compute language that compiles to CUDA, ROCm, Metal, Vulkan and WebGPU. Kernels are written once in Rust and run on any supported backend.
Model interop	burn‑onnx	Imports ONNX models, turning them into native Burn Rust code.
	burn‑store	Loads and saves weights (PyTorch .pt, Safetensors, etc.).
Domain libraries	burn‑vision, burn‑rl, burn‑dataset	Vision ops, reinforcement‑learning helpers, and dataset utilities.
Pre‑trained models	models	A collection of ready‑to‑run examples built with Burn.
Tooling	burn‑bench	Benchmarks for comparing backend performance over time.

Supported back‑ends – Burn can run on many devices through a plug‑in architecture:

• GPU – CUDA, ROCm, Metal, Vulkan, WebGPU (covers Nvidia, AMD, Apple, Intel, Qualcomm and browsers).
• CPU – Two options:
  • CubeCL CPU – JIT‑compiled kernels via LLVM (supports x86 and ARM).
  • Flex – Pure‑Rust eager backend that works on no_std, WebAssembly and embedded targets.
• Back‑ends are interchangeable at runtime via the Device abstraction; you can enable several in the same binary and pick the device per‑tensor.

Special features

• Autodiff – Back‑propagation is a decorator that can be turned on per device. Tensors record a computation graph only when you request it, keeping inference lightweight.
• Kernel fusion – Automatically merges consecutive operations into a single GPU kernel, reducing launch overhead.
• Remote execution (beta) – A client‑server model lets you off‑load tensor work to another machine (or cloud) using encrypted Iroh or simple WebSocket transport.
• Training dashboard – A terminal UI (built with Ratatui) shows live metrics, lets you pause training safely, and works without external services.
• ONNX & PyTorch weight import – Bring existing models into Burn without rewriting them.
• WebAssembly inference – Flex (CPU) and WGPU back‑ends let you run models directly in the browser.

Typical workflow

1. Define a model using Burn’s Module trait, just like you would in PyTorch.
2. Create a Device (e.g., Device::cuda(0).autodiff() or Device::wgpu(Default::default()).autodiff()).
3. Run training – tensors automatically record gradients; call backward() to compute them.
4. Save weights with burn-store or export to ONNX for interoperability.
5. Deploy – Switch the Device to a CPU, WebGPU, or remote backend without changing model code.

Who should use it

• Researchers who want Rust’s safety and speed but still need a rapid edit‑run loop.
• Engineers building production pipelines that must run on heterogeneous hardware (cloud GPUs, edge CPUs, browsers) from a single code base.
• Contributors interested in expanding the Rust AI ecosystem (new back‑ends, domain crates, model zoo).

Getting started – The crate is published on crates.io (burn = "<latest>"). Documentation lives at https://burn.dev/docs/burn, and the “Burn Book” provides step‑by‑step tutorials for tensors, autodiff, training loops, and deployment.

---

All information above is taken directly from the repository’s README; no additional features have been inferred.