tracel-ai/burn

These crates are not maintained by Tracel, but they are part of the same Rust AI story. Anything that helps you load data, build environments, or ship models belongs here. Built something that fits? Open a PR to add it!

| Category | Crate | Description | | -------------------------- | --------------------------------------------------------------- | ----------------------------------------------------------------- | | Data & loading | polars | Fast DataFrames for tabular data | | | arrow-rs | Apache Arrow columnar memory format | | | image | Image decoding, encoding and processing | | | hf-hub | Download models and datasets from the Hugging Face Hub | | Tokenization & NLP | tokenizers | Fast, production-ready tokenizers | | | rust-bert | Ready-to-use NLP pipelines and transformer models | | Numerical & linear algebra | ndarray | N-dimensional arrays | | | nalgebra | Linear algebra | | Classical ML | linfa | Classical ML toolkit, in the spirit of scikit-learn | | | smartcore | Classical ML algorithms, no BLAS/LAPACK required | | Inference & runtimes | candle | Minimalist ML framework with a focus on LLM inference | | | mistral.rs | Fast, multimodal LLM inference engine | | | ort | ONNX Runtime bindings for hardware-accelerated inference | | | tract | Pure-Rust inference for ONNX and NNEF models | | | wonnx | 100% Rust, WebGPU-accelerated ONNX runtime for native and the web | | LLM apps & RAG | rig | Build modular LLM applications and agents | | | langchain-rust | LangChain-style chain orchestration | | Embeddings & vector search | fastembed | Generate text embeddings and rerank locally | | | qdrant | Vector search engine, written in Rust | | | lancedb | Embedded, developer-friendly vector database | | Computer vision | kornia-rs | Low-level 3D computer vision library | | Simulation & environments | rapier | Physics engine for robotics and RL environments | | Visualization | rerun | Multimodal data and CV/robotics visualization | | | plotters | Plotting and charting |

In application code, autodiff is runtime context carried by tensors. Devices provide the default context for newly created tensors, and each tensor can later enable or remove autodiff independently. Enabling autodiff permits graph recording; it does not by itself make a tensor retain gradients.

Internally, Burn implements this by decorating a concrete backend, so autodiff cannot execute by itself. Enable it on a device before creating tensors or initializing a model. With the autodiff and wgpu features enabled:

use burn::tensor::{Device, Distribution, Tensor};
fn main() {
    let device = Device::wgpu(Default::default()).autodiff();
let x: Tensor<2> = Tensor::random([32, 32], Distribution::Default, &device);
    let y: Tensor<2> = Tensor::random([32, 32], Distribution::Default, &device).require_grad();
let tmp = x.clone() + y.clone();
    let tmp = tmp.matmul(x);
    let tmp = tmp.exp();
let grads = tmp.backward();
    let y_grad = y.grad(&grads).unwrap();
    println!("{y_grad}");
}

backward() checks graph participation at runtime. Enable autodiff before the forward pass and call require_grad() on source leaves whose gradients you need. is_autodiff(), is_tracked(), and is_require_grad() inspect autodiff association, graph participation, and gradient retention respectively. See the autodiff guide.

See the Autodiff Backend README for more details.

This backend decorator enhances a backend with kernel fusion, provided that the inner backend supports it. Note that you can compose this backend with other backend decorators such as Autodiff. All first-party accelerated backends (like WGPU and CUDA) use Fusion by default (burn/fusion feature flag), so you typically don't need to apply it manually.

#[cfg(not(feature = "fusion"))]
pub type Cube = burn_cubecl::CubeBackend;
[cfg(feature = "fusion")]
pub type Cube = burn_fusion::Fusion<burn_cubecl::CubeBackend>;

Device::autodiff().gradient_checkpointing() enables the balanced gradient-checkpointing strategy, which trades recomputation for reduced activation storage during training.

See the Fusion Backend README for more details.

Remote execution has a client and a server. The server's devices select the compute backend; clients use a remote Device with the same tensor API. Iroh, the default transport, reaches a server across any network, authenticated and encrypted; see the server example and the distributed computing guide. On a trusted network, WebSocket is the simplest setup: enable remote-server, remote-websocket, and cuda on the server, and remote-websocket plus autodiff on the client:

use burn::remote::RemoteHost;
use burn::server::{RemoteServer, ServeError, WebSocketTransport};
use burn::tensor::{Device, Distribution, Tensor};
fn main_server() -> Result<(), ServeError> {
    RemoteServer::new([Device::cuda(0)]).serve(WebSocketTransport::new(3000))
}
fn main_client() -> Result<(), burn::remote::ConnectError> {
    let host = RemoteHost::websocket("ws://localhost:3000");
    let device = Device::remote_options(&host).init()?.autodiff();
    let tensor_gpu = Tensor::<2>::random([3, 3], Distribution::Default, &device);
    Ok(())
}

As you can see in the previous video (click on the picture!), a new terminal UI dashboard based on the Ratatui crate allows users to follow their training with ease without having to connect to any external application.

You can visualize your training and validation metrics updating in real-time and analyze the lifelong progression or recent history of any registered metrics using only the arrow keys. Break from the training loop without crashing, allowing potential checkpoints to be fully written or important pieces of code to complete without interruption 🛡

Burn supports importing ONNX (Open Neural Network Exchange) models through the burn-onnx crate, allowing you to easily port models from TensorFlow or PyTorch to Burn. The ONNX model is converted into Rust code that uses Burn's native APIs, enabling the imported model to run on any Burn backend (CPU, GPU, WebAssembly) and benefit from all of Burn's optimizations like automatic kernel fusion.

Our ONNX support is further described in this section of the Burn Book 🔥.

Note: This crate is in active development and currently supports a

limited set of ONNX operators.

You can load weights from PyTorch or Safetensors formats directly into your Burn-defined models. This makes it easy to reuse existing models while benefiting from Burn's performance and deployment features.

Learn more in the Saving & Loading Models section of the Burn Book.

Several of our backends can run in WebAssembly environments: Flex for CPU execution, and WGPU for GPU acceleration via WebGPU. This means that you can run inference directly within a browser. We provide several examples of this:

• MNIST where you can draw digits and a small convnet tries to

• Image Classification

Burn's core components support no_std. This means it can run in bare metal environment such as embedded devices without an operating system.

As of now, only the Flex backend can be used in a _no_std_ environment.

To begin working effectively with Burn, it is crucial to understand its key components and philosophy. This is why we highly recommend new users to read the first sections of The Burn Book 🔥. It provides detailed examples and explanations covering every facet of the framework, including building blocks like tensors, modules, and optimizers, all the way to advanced usage, like coding your own GPU kernels.

The project is constantly evolving, and we try as much as possible to keep the book up to date

with new additions. However, we might miss some details sometimes, so if you see something weird,

let us know! We also gladly accept Pull Requests 😄

Let's start with a code snippet that shows how intuitive the framework is to use! In the following, we declare a neural network module with some parameters along with its forward pass.

use burn::nn;
use burn::module::Module;
use burn::tensor::Tensor;
[derive(Module, Debug)]
pub struct PositionWiseFeedForward {
    linear_inner: nn::Linear,
    linear_outer: nn::Linear,
    dropout: nn::Dropout,
    gelu: nn::Gelu,
}
impl PositionWiseFeedForward {
    pub fn forward(&self, input: Tensor) -> Tensor {
        let x = self.linear_inner.forward(input);
        let x = self.gelu.forward(x);
        let x = self.dropout.forward(x);
self.linear_outer.forward(x)
    }
}

We have a somewhat large amount of examples in the repository that shows how to use the framework in different scenarios.

Following the book:

• Basic Workflow : Creates a custom CNN Module to train on the MNIST dataset

• Custom Training Loop : Implements a basic training loop instead

• Custom WGPU Kernel : Learn how to create your own custom

Additional examples:

• Custom CSV Dataset : Implements a dataset to parse CSV data for a

• Regression : Trains a simple MLP on the California Housing dataset

• [Custom Image Dataset](https://github.com/tracel-ai/burn/tree/main/examples/custom