Skip to main content

burn/
lib.rs

1#![cfg_attr(not(feature = "std"), no_std)]
2#![warn(missing_docs)]
3#![cfg_attr(docsrs, feature(doc_cfg))]
4
5//! # Burn
6//!
7//! Burn is a deep learning framework written in Rust. It covers tensors, automatic
8//! differentiation, neural network modules, optimizers, training and model storage, and runs
9//! the same model code on GPUs (CUDA, ROCm, Metal, Vulkan, WebGPU), CPUs and WebAssembly.
10//!
11//! ## Quick start
12//!
13//! Burn ships no execution backend by default. Enable one or more with Cargo features:
14//!
15//! ```toml
16//! [dependencies]
17//! burn = { version = "0.22", features = ["wgpu"] }
18//! ```
19//!
20//! Models are plain structs that derive [`Module`](module::Module). Tensors carry their rank in
21//! the type and their backend in their [`Device`](tensor::Device), so model code has no backend
22//! type parameter:
23//!
24//! ```rust,no_run
25//! use burn::nn::{Linear, LinearConfig, Relu};
26//! use burn::prelude::*;
27//!
28//! #[derive(Module, Debug)]
29//! struct Mlp {
30//!     hidden: Linear,
31//!     activation: Relu,
32//!     output: Linear,
33//! }
34//!
35//! impl Mlp {
36//!     fn new(device: &Device) -> Self {
37//!         Self {
38//!             hidden: LinearConfig::new(784, 128).init(device),
39//!             activation: Relu::new(),
40//!             output: LinearConfig::new(128, 10).init(device),
41//!         }
42//!     }
43//!
44//!     fn forward(&self, input: Tensor<2>) -> Tensor<2> {
45//!         let x = self.activation.forward(self.hidden.forward(input));
46//!         self.output.forward(x)
47//!     }
48//! }
49//!
50//! // An enabled backend in priority order (GPUs before CPUs), unless `BURN_DEVICE` names one.
51//! // `Device::wgpu(..)`, `Device::cuda(0)`, ... pick one explicitly.
52//! let device = Device::default();
53//! let model = Mlp::new(&device);
54//! let logits = model.forward(Tensor::zeros([32, 784], &device));
55//! ```
56//!
57//! The [Burn Book](https://burn.dev/books/burn/) walks through a full training workflow.
58//!
59//! ## Crate map
60//!
61//! - [`tensor`]: [`Tensor`], [`Device`](tensor::Device), dtypes and tensor
62//!   operations.
63//! - [`module`] and [`nn`]: the [`Module`](module::Module) trait and neural network layers.
64//! - [`config`]: serializable configuration structs with `#[derive(Config)]`.
65//! - [`optim`], [`lr_scheduler`], [`grad_clipping`]: optimizers and training utilities.
66//! - [`data`]: datasets, transformations and data loaders.
67//! - [`store`]: saving and loading weights in burnpack, SafeTensors and PyTorch formats.
68//! - `train`: the `Learner`, metrics and the training dashboard (`train` feature).
69//! - `vision`, `signal`, `linalg`: domain-specific tensor operations (features of the same
70//!   names).
71//! - `remote` and `server`: run tensors on devices hosted by another machine (`remote` and
72//!   `remote-server` features).
73//! - [`prelude`]: the types most programs import.
74//!
75//! ## Backends
76//!
77//! Every enabled backend is available at runtime through a `Device` constructor, and several can
78//! be used side by side:
79//!
80//! | Backend                 | Feature  | Device                               |
81//! | ----------------------- | -------- | ------------------------------------ |
82//! | CUDA                    | `cuda`   | `Device::cuda(0)`                    |
83//! | ROCm                    | `rocm`   | `Device::rocm(0)`                    |
84//! | wgpu (any graphics API) | `wgpu`   | `Device::wgpu(Default::default())`   |
85//! | Metal                   | `metal`  | `Device::metal(Default::default())`  |
86//! | Vulkan                  | `vulkan` | `Device::vulkan(Default::default())` |
87//! | WebGPU                  | `webgpu` | `Device::webgpu(Default::default())` |
88//! | CubeCL CPU              | `cpu`    | `Device::cpu()`                      |
89//! | Flex (pure Rust CPU)    | `flex`   | `Device::flex()`                     |
90//!
91//! Autodiff and kernel fusion are decorators over these backends: `device.autodiff()` enables
92//! gradients for tensors created on a device, and the CubeCL backends fuse operations by default.
93//! NdArray (`ndarray`) and LibTorch (`tch`) are deprecated.
94//!
95//! ## Quantization
96//!
97//! Burn supports post-training quantization of weights and activations, per tensor or per
98//! block, to 8, 4 and 2-bit integers and to FP8 and FP4 formats on supported backends.
99//! Quantization-aware training is not supported yet. See the
100//! [quantization chapter](https://burn.dev/books/burn/performance/quantization.html).
101//!
102//! ## Feature Flags
103//!
104//! The following feature flags are available.
105//! Default features include `std` and `optim` (and therefore `autodiff`), but no execution backend.
106//! Select a backend explicitly, for example `features = ["wgpu"]` or `["flex"]`.
107//! Specialized operations are also opt-in, for example `features = ["flex", "signal"]`.
108//! Backend-free builds can define tensor/model APIs without installing an execution backend.
109//! `Device::default()` panics if no execution backend is available; graph capture remains
110//! available through `Device::capture()` with the `capture` feature.
111//!
112//! - Training
113//!   - `train`: Enables features `dataset` and `optim` and provides a training environment
114//!   - `optim`: Enables optimizers and learning rate schedulers (implies `autodiff`)
115//!   - `rl`: Enables reinforcement learning utilities
116//!   - `tui`: Includes Text UI with progress bar and plots (requires `train`)
117//!   - `metrics`: Includes system info metrics (CPU/GPU usage, etc.) (requires `train`)
118//! - Dataset
119//!   - `dataset`: Includes a datasets library
120//!   - `audio`: Enables audio datasets (SpeechCommandsDataset)
121//!   - `sqlite`: Stores datasets in an SQLite database, backed by [Turso](https://turso.tech/)
122//!   - `sqlite-bundled`: Deprecated alias for `sqlite`
123//!   - `vision`: Enables vision datasets (MnistDataset) and the `burn-vision` ops module
124//! - Backends
125//!   - `wgpu`: Makes available the WGPU backend, on whichever graphics API the platform provides
126//!   - `webgpu`: Adds `Device::webgpu`, pinned to WebGPU (implies `wgpu`)
127//!   - `vulkan`: Adds `Device::vulkan`, pinned to Vulkan (implies `wgpu`)
128//!   - `metal`: Adds `Device::metal`, pinned to Metal with native MSL (implies `wgpu`)
129//!   - `cuda`: Makes available the CUDA backend
130//!   - `rocm`: Makes available the ROCm backend
131//!   - `cpu`: Makes available the CubeCL CPU backend
132//!   - `tch`: Makes available the LibTorch backend (deprecated - use a CubeCL backend instead)
133//!   - `flex`: Makes available the Flex backend (pure-Rust CPU, std/no_std/WASM)
134//!   - `ndarray`: Makes available the NdArray backend (deprecated - use `flex` instead)
135//! - Backend specifications
136//!   - `simd`: Enable SIMD kernels in the Flex and NdArray backends
137//!   - `rayon`: Enable multi-threaded execution in the Flex and NdArray backends
138//!   - `accelerate`, `blas-netlib`, `openblas`, `openblas-system`: BLAS providers for the NdArray
139//!     backend
140//!   - `autotune`: Enable running benchmarks to select the best kernel in backends that support it.
141//!   - `autotune-checks`: Check that every autotune candidate produces the same output (debugging).
142//!   - `x86-v4`: Enable AVX-512 matmul kernels in the Flex backend.
143//!   - `apple-amx`: Enable the experimental Apple AMX matmul kernels in the Flex backend.
144//!   - `template`: Enable hand-written, non-JIT custom kernels in the CubeCL backends.
145//!   - `fusion`: Enable operation fusion in backends that support it.
146//!   - `tracing`: Enable diagnostic tracing in the selected backends (disabled by default).
147//! - Backend decorators
148//!   - `autodiff`: Makes available the Autodiff backend
149//! - Model Storage
150//!   - `store`: Enables the `burn-store` snapshot tooling and burnpack stores; with `std`, this
151//!     also includes SafeTensors
152//!   - `safetensors`: Enables SafeTensors import and export in `no_std` builds (implies `store`)
153//!   - `pytorch`: Enables PyTorch checkpoint import (implies `store`)
154//! - Others:
155//!   - `std`: Activates the standard library (deactivate for no_std)
156//!   - `linalg`: Enables linear algebra operations
157//!   - `capture`: Makes the non-executing graph capture backend available.
158//!   - `ir`: Makes Burn's operation intermediate representation available.
159//!   - `cubecl`: Re-exports CubeCL as `burn::cubecl` for writing custom kernels.
160//!   - `signal`: Enables signal processing operations from `burn-signal`.
161//!   - `extension`: Enables the backend extension API, including `Tensor::from_primitive`.
162//!   - `remote`: Enables remote devices over Iroh; `remote-websocket` adds the WebSocket transport.
163//!   - `remote-server`: Enables the remote server (implies `remote`).
164//!   - `network`: Enables network utilities (currently, only a file downloader with progress bar)
165//!
166//! You can also check the details in sub-crates [`burn-core`](https://docs.rs/burn-core) and [`burn-train`](https://docs.rs/burn-train).
167//!
168//! ### Backend tracing
169//!
170//! Add `"tracing"` to the features of your `burn` dependency to compile backend instrumentation,
171//! including autodiff and fusion spans. When depending directly on `burn-autodiff` or
172//! `burn-fusion`, enable their `tracing` feature instead. These spans are opt-in: configuring a
173//! tracing subscriber alone does not enable them. Configure your subscriber to include the
174//! `trace` level to observe tensor operation spans.
175//!
176//! The feature propagates to enabled backends without selecting an additional backend. Normal
177//! training logs remain available without this feature.
178
179pub use burn_core::*;
180
181/// Linear algebra operations.
182#[cfg(feature = "linalg")]
183pub mod linalg {
184    pub use burn_linalg::*;
185}
186
187/// Core module infrastructure and neural-network initializers.
188pub mod module {
189    pub use burn_core::module::*;
190    pub use burn_nn::Initializer;
191}
192
193/// Tensor types and compatibility re-exports.
194pub mod tensor {
195    pub use burn_core::tensor::*;
196
197    /// Compatibility path for signal processing operations.
198    #[cfg(feature = "signal")]
199    pub mod signal {
200        pub use burn_signal::*;
201    }
202
203    /// Compatibility path for linear algebra operations.
204    #[cfg(feature = "linalg")]
205    pub mod linalg {
206        pub use burn_linalg::*;
207    }
208}
209
210/// Train module
211#[cfg(feature = "train")]
212pub mod train {
213    pub use burn_train::*;
214}
215
216/// Module for reinforcement learning.
217#[cfg(feature = "rl")]
218pub mod rl {
219    pub use burn_rl::*;
220}
221
222#[cfg(feature = "remote")]
223pub use burn_core::tensor::remote;
224#[cfg(feature = "remote-server")]
225pub use burn_core::tensor::server;
226
227/// Model storage and serialization: the non-generic record system (always available), plus,
228/// with the `store` feature, the snapshot tooling and burnpack stores. The `safetensors` and
229/// `pytorch` features add those importers.
230pub mod store {
231    pub use burn_core::store::*;
232    #[cfg(feature = "store")]
233    pub use burn_store::*;
234}
235
236/// Neural network module.
237pub mod nn {
238    pub use burn_nn::*;
239}
240
241pub use burn_std::config::{BurnConfig, config as runtime_config};
242
243#[cfg(all(test, feature = "capture"))]
244mod capture_tests {
245    use crate::{module::Module, nn::BatchNormConfig, tensor::Device};
246
247    #[test]
248    fn capture_feature_exposes_the_user_facing_device_api() {
249        let device = Device::capture();
250        let captured = device
251            .capture_scope(|scope| scope.complete([], []))
252            .unwrap();
253
254        assert!(captured.graph.operations.is_empty());
255    }
256
257    #[test]
258    fn shared_running_state_moves_across_capture_scopes() {
259        let module = BatchNormConfig::new(3).init(&Device::default());
260        let first_device = Device::capture();
261        let second_device = Device::capture();
262
263        let first = first_device
264            .capture_scope(|scope| {
265                let _module = module.clone().to_device(&first_device);
266                scope.complete([], [])
267            })
268            .unwrap();
269        let second = second_device
270            .capture_scope(|scope| {
271                let _module = module.clone().to_device(&second_device);
272                scope.complete([], [])
273            })
274            .unwrap();
275
276        assert_eq!(first.values.len(), 4);
277        assert_eq!(second.values.len(), 4);
278    }
279}
280
281/// Optimizers module.
282#[cfg(feature = "optim")]
283pub mod optim {
284    pub use burn_optim::*;
285}
286
287// For backward compat, `burn::lr_scheduler::*`
288/// Learning rate scheduler module.
289#[cfg(all(feature = "optim", feature = "std"))]
290pub mod lr_scheduler {
291    pub use burn_optim::lr_scheduler::*;
292}
293// For backward compat, `burn::grad_clipping::*`
294/// Gradient clipping module.
295#[cfg(feature = "optim")]
296pub mod grad_clipping {
297    pub use burn_optim::grad_clipping::*;
298}
299
300/// CubeCL module re-export.
301#[cfg(feature = "cubecl")]
302pub mod cubecl {
303    pub use cubecl::*;
304}
305
306#[cfg(feature = "vision")]
307/// Vision module.
308pub mod vision {
309    pub use burn_vision::*;
310}
311
312#[cfg(feature = "signal")]
313/// Signal processing module.
314pub mod signal {
315    pub use burn_signal::*;
316}
317
318pub mod prelude {
319    //! Structs and macros used by most projects. Add `use
320    //! burn::prelude::*` to your code to quickly get started with
321    //! Burn.
322    pub use burn_core::prelude::*;
323
324    pub use crate::nn;
325}