burn

Burn is a tensor library and deep learning framework written in Rust, aimed at numerical computing, training and inference. Its stated goal is to avoid the usual split between a Python training stack and a separate deployment engine: the same model code can be trained and then run in production, which the project presents as useful for on-device personalization and federated learning. Design and ergonomics - PyTorch-like API with dynamic shapes and graphs, while JIT-compiling streams of tensor operations and performing automatic kernel fusion. - Incremental compilation is emphasized: the README claims model code changes recompile in under 5 seconds even in release mode. - Backends can be swapped without changing model code; multiple backends can coexist in one application and the device is chosen at runtime via `Device`. Backends - GPU: CUDA, ROCm, Metal, Vulkan, WebGPU, and LibTorch (deprecated as of 0.22.0). Coverage varies by vendor: Nvidia (CUDA, Vulkan, WebGPU, LibTorch), AMD (ROCm, Vulkan, WebGPU, LibTorch), Apple (Metal, WebGPU, LibTorch), Intel and Qualcomm (Vulkan, WebGPU), Wasm (WebGPU). - CPU: CubeCL CPU backend, Flex, and LibTorch for x86 and Arm; Flex also covers Wasm and no_std. - Backend decorators: Autodiff adds backpropagation to any backend; Fusion adds kernel fusion where supported and is enabled by default on first-party accelerated backends; Remote (beta) supports client/server execution over Iroh or WebSocket for distributed computation. Training and inference - A terminal training dashboard built on Ratatui shows training and validation metrics in real time, with navigation via arrow keys and the ability to break from the training loop without crashing. - ONNX models can be imported through burn-onnx and converted into native Rust code using Burn APIs, so they run on any Burn backend; the README notes this crate is in active development with a limited operator set. - Weights from PyTorch and Safetensors formats can be loaded into Burn-defined models. - Inference can run in the browser via WebAssembly using Flex (CPU) or WGPU (WebGPU), with MNIST and image-classification demos. - Core components support no_std for bare-metal embedded environments; currently only the Flex backend works in no_std. Ecosystem - CubeCL: GPU compute language and compiler behind the accelerated backends, usable standalone. - burn-onnx: ONNX import; burn-store: weight save/load and PyTorch/Safetensors import. - burn-vision, burn-rl, burn-dataset for vision, reinforcement learning and datasets. - models repository with pre-trained models and examples; burn-bench for benchmarking backends over time. - A community crate list covers data loading (polars, arrow-rs, image, hf-hub), tokenization and NLP (tokenizers, rust-bert), numerical libraries (ndarray, nalgebra), classical ML (linfa, smartcore), inference runtimes (candle, mistral.rs, ort, tract, wonnx), LLM/RAG tooling (rig, langchain-rust), embeddings and vector search (fastembed, qdrant, lancedb), computer vision (kornia-rs), simulation (rapier) and visualization (rerun, plotters). Getting started - The Burn Book is the primary documentation, covering tensors, modules, optimizers and custom GPU kernels. - Examples include a basic MNIST workflow, custom training loop, custom WGPU kernel, CSV and image datasets, regression, custom renderer, browser inference demos, PyTorch weight import, text classification and text generation, and WGAN on MNIST. - Licensing is MIT/Apache-2.0.