burn Documentation

File: README.md Copy Chapter

Training and inference usually live in separate worlds. Models are typically trained in Python then
exported to an open format like ONNX or optimized for production engines like vLLM, ONNX Runtime, or
TensorRT. This export step is often brittle and lossy, ruling out complex architectures and advanced
deployment use cases.

Burn unifies the two. By executing multi-platform tensor operations via a single, unified API, the
exact code used for training is the exact code that runs in production. This makes workloads like
on-device personalization and federated learning straightforward, while enabling teams to go from
prototype to deployment in a single codebase.

Burn preserves the intuitive ergonomics of PyTorch, with dynamic shapes and graphs, but JIT-compiles
streams of tensor operations, performing automatic kernel fusion. You get the flexibility of dynamic
graphs without the performance drop.

Rust for Research? Copy Chapter

Rust used to be a tough sell for research: long compilation times disrupted the fast
edit-compile-run loop that draws researchers to Python. Burn changes this paradigm. Designed around
incremental compilation, modifying model code recompiles in under 5 seconds, even in release mode.
This delivers a Python-like feedback loop with the speed and safety of Rust.

Ecosystem Copy Chapter

Burn is the core of a growing, fully open-source Rust AI ecosystem. You are not adopting a single
library, you are joining a stack that spans GPU compute, model interop and domain toolkits, with
plenty of room to help shape what comes next.

Burn's CubeCL backends (CUDA, ROCm, Metal, Vulkan, WebGPU,
CPU) compose with autodiff, fusion and remote-execution decorators, while external and simpler
backends (LibTorch and pure-Rust CPU/no_std) compose with autodiff only. See
Supported Backends below for the full matrix.

Every project here is open-source and actively developed. Want to help build the Rust AI ecosystem?
The good first issues are a great place to start,
and the Contributing guide will get you set up.

Community crates 🌱

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!

Backend Copy Chapter

Burn strives to be as fast as possible on as many hardwares as possible, with robust
implementations. We believe this flexibility is crucial for modern needs where you may train your
models in the cloud, then deploy on customer hardwares, which vary from user to user.

Supported Backends Copy Section

Most backends support all operating systems, so we don't mention them in the tables below.

GPU Backends:

CPU Backends:

Compared to other frameworks, Burn has a very different approach to supporting many backends. By
design, most code is generic over the Backend trait, which allows us to build Burn with swappable
backends. This makes composing backend possible, augmenting them with additional functionalities
such as autodifferentiation and automatic kernel fusion.

Autodiff: Backend decorator that brings backpropagation to any backend 🔄

Contrary to the aforementioned backends, Autodiff is actually a backend decorator. This means that
it cannot exist by itself; it must encapsulate another backend.

The simple act of wrapping a base backend with Autodiff transparently equips it with
autodifferentiation support, making it possible to call backward on your model.

use burn::backend::{Autodiff, Wgpu};
use burn::tensor::{Distribution, Tensor};

fn main() {
    type Backend = Autodiff<Wgpu>;

    let device = Default::default();

    let x: Tensor = Tensor::random([32, 32], Distribution::Default, &device);
    let y: Tensor = 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}");
}

Of note, it is impossible to make the mistake of calling backward on a model that runs on a backend
that does not support autodiff (for inference), as this method is only offered by an Autodiff
backend.

See the Autodiff Backend README for more details.

Fusion: Backend decorator that brings kernel fusion to all first-party backends

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 Cuda<F = f32, I = i32> = CubeBackend<CudaRuntime, F, I, u8>;

#[cfg(feature = "fusion")]
pub type Cuda<F = f32, I = i32> = burn_fusion::Fusion<CubeBackend<CudaRuntime, F, I, u8>>;

Of note, we plan to implement automatic gradient checkpointing based on compute bound and memory
bound operations, which will work gracefully with the fusion backend to make your code run even
faster during training, see this issue.

See the Fusion Backend README for more details.

Remote (Beta): Backend decorator for remote backend execution, useful for distributed computations

That backend has two parts, one client and one server. The client sends tensor operations over the
network to a remote compute backend. You can use any first-party backend as server in a single line
of code:

fn main_server() {
    // Start a server on port 3000.
    burn::server::start::(Default::default(), 3000);
}

fn main_client() {
    // Create a client that communicate with the server on port 3000.
    use burn::backend::{Autodiff, RemoteBackend};

    type Backend = Autodiff<RemoteDevice>;

    let device = RemoteDevice::new("ws://localhost:3000");
    let tensor_gpu =
        Tensor::::random([3, 3], Distribution::Default, &device);
}

Training & Inference Copy Chapter

The whole deep learning workflow is made easy with Burn, as you can monitor your training progress
with an ergonomic dashboard, and run inference everywhere from embedded devices to large GPU
clusters.

Burn was built from the ground up with training and inference in mind. It's also worth noting how
Burn, in comparison to frameworks like PyTorch, simplifies the transition from training to
deployment, eliminating the need for code changes.

Click on the following sections to expand 👇

Training Dashboard 📈

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 🛡

ONNX Support 🐫

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.

Importing PyTorch or Safetensors Models 🚚

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.

Inference in the Browser 🌐

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
  find which one it is! 2️⃣ 7️⃣ 😰
• Image Classification
  where you can upload images and classify them! 🌄

Embedded: no_std support ⚙️

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.

Benchmarks Copy Section

To evaluate performance across different backends and track improvements over time, we provide a
dedicated benchmarking suite.

Run and compare benchmarks using burn-bench.

⚠️ Warning When using one of the wgpu backends, you may encounter compilation errors related
to recursive type evaluation. This is due to complex type nesting within the wgpu dependency
chain. To resolve this issue, add the following line at the top of your main.rs or lib.rs
file:

rust Copy

#![recursion_limit = "256"]

The default recursion limit (128) is often just below the required depth (typically 130-150) due
to deeply nested associated types and trait bounds.

Getting Started Copy Chapter

Just heard of Burn? You are at the right place! Just continue reading this section and we hope you
can get on board really quickly.

The Burn Book 🔥

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 😄

Examples 🙏

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::backend::Backend;

#[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<const D: usize>(&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
  and use for inference.
• Custom Training Loop : Implements a basic training loop instead
  of using the Learner.
• Custom WGPU Kernel : Learn how to create your own custom
  operation with the WGPU backend.

Additional examples:

• Custom CSV Dataset : Implements a dataset to parse CSV data for a
  regression task.
• Regression : Trains a simple MLP on the California Housing dataset
  to predict the median house value for a district.
• Custom Image Dataset : Trains a simple CNN on custom image
  dataset following a simple folder structure.
• Custom Renderer : Implements a custom renderer to display the
  Learner progress.
• Image Classification Web : Image classification web browser
  demo using Burn, WGPU and WebAssembly.
• MNIST Inference on Web : An interactive MNIST inference demo in
  the browser. The demo is available online.
• MNIST Training : Demonstrates how to train a custom Module (MLP) with the
  Learner configured to log metrics and keep training checkpoints.
• PyTorch Import Inference : Imports a PyTorch model pre-trained
  on MNIST to perform inference on a sample image with Burn.
• Text Classification : Trains a text classification transformer
  model on the AG News or DbPedia dataset. The trained model can then be used to classify a text
  sample.
• Text Generation : Trains a text generation transformer model on the
  DbPedia dataset.
• Wasserstein GAN MNIST : Trains a WGAN model to generate new handwritten digits
  based on MNIST.

For more practical insights, you can clone the repository and run any of them directly on your
computer!

Pre-trained Models 🤖

We keep an updated and curated list of models and examples built with Burn, see the
tracel-ai/models repository for more details.

Don't see the model you want? Don't hesitate to open an issue, and we may prioritize it. Built a
model using Burn and want to share it? You can also open a Pull Request and add your model under the
community section!

Why use Rust for AI? 🦀

Deep Learning is a special form of software where you need very high level abstractions as well as
extremely fast execution time. Rust is the perfect candidate for that use case since it provides
zero-cost abstractions to easily create neural network modules, and fine-grained control over memory
to optimize every detail. To this day, the mainstream solution has been to offer APIs in Python but
rely on bindings to low-level languages such as C/C++. This reduces portability, increases
complexity and creates friction between researchers and engineers. Rust's approach to abstractions
is versatile enough to tackle this two-language dichotomy, and Cargo makes it easy to build, test
and deploy from any environment, which is usually a pain in Python.

Rust's AI ecosystem is young, but it is real and growing quickly. Foundational pieces are already
here: Burn and CubeCL for training and compute,
candle for inference, Hugging Face's tokenizers and
safetensors, and polars and ndarray for data. Betting on Rust today means betting on a stack
that is growing, and one where contributors still shape the direction. The pieces that don't exist
yet are opportunities rather than dead-ends (see Contributing).

Rust is also what makes one-stack-everywhere possible: a single self-contained binary with no Python
runtime to ship, running from servers down to no_std embedded targets.

Deprecation NoteSince 0.14.0, the internal structure for tensor data has changed. The
previous Data struct was deprecated and officially removed since 0.17.0 in favor of the new
TensorData struct, which allows for more flexibility by storing the underlying data as bytes and
keeping the data type as a field. If you are using Data in your code, make sure to switch to
TensorData.

Loading Model Records From Previous Versions ⚠️

In the event that you are trying to load a model record saved in a version older than 0.14.0, make
sure to use a compatible version (0.14, 0.15 or 0.16) with the record-backward-compat
feature flag.

features = [..., "record-backward-compat"]

Otherwise, the record won't be deserialized correctly and you will get an error message. This error
will also point you to the backward compatible feature flag.

The backward compatibility was maintained for deserialization when loading records. Therefore, as
soon as you have saved the record again it will be saved according to the new structure and you can
upgrade back to the current version

Please note that binary formats are not backward compatible. Thus, you will need to load your record
in a previous version and save it in any of the other self-describing record format (e.g., using the
NamedMpkFileRecorder) before using a compatible version (as described) with the
record-backward-compat feature flag.

Community Copy Chapter

If you are excited about the project, don't hesitate to join our
Discord! We try to be as welcoming as possible to everybody from
any background. You can ask your questions and share what you built with the community!

Contributing Copy Section

Before contributing, please read the Contributing Guidelines and our
Code of Conduct. The Contributor Book
covers architecture, environment setup, and guides for common tasks.

Status Copy Chapter

Burn is currently in active development, and there will be breaking changes. While any resulting
issues are likely to be easy to fix, there are no guarantees at this stage.

License Copy Chapter

Burn is distributed under the terms of both the MIT license and the Apache License (Version 2.0).
See LICENSE-APACHE and LICENSE-MIT for details. Opening a pull
request is assumed to signal agreement with these licensing terms.

File: burn-book/src/building-blocks/README.md Copy Chapter

Building Blocks

In this section, we'll guide you through the core elements that make up Burn. We'll walk you through
the key components that serve as the building blocks of the framework and your future projects.

As you explore Burn, you might notice that we occasionally draw comparisons to PyTorch. We believe
it can provide a smoother learning curve and help you grasp the nuances more effectively.

File: burn-book/src/basic-workflow/README.md Copy Chapter

Guide

This guide will walk you through the process of creating a custom model built with Burn. We will
train a simple convolutional neural network model on the MNIST dataset and prepare it for inference.

For clarity, we sometimes omit imports in our code snippets. For more details, please refer to the
corresponding code in the examples/guide directory.
We reproduce this example in a step-by-step fashion, from dataset creation to modeling and training
in the following sections. It is recommended to use the capabilities of your IDE or text editor to
automatically add the missing imports as you add the code snippets to your code.

Be sure to checkout the git branch corresponding to the version of Burn you are using to follow
this guide.

The current version of Burn is 0.21 and the corresponding branch to checkout is main.

The code for this demo can be executed from Burn's base directory using the command:

cargo run --release --example guide

Key Learnings Copy Chapter

• Creating a project
• Creating neural network models
• Importing and preparing datasets
• Training models on data
• Choosing a backend
• Using a model for inference

File: burn-book/src/advanced/backend-extension/README.md Copy Chapter

Backend Extension

Burn aims to be the most flexible deep learning framework. While it's crucial to maintain
compatibility with a wide variety of backends, Burn provides the ability to extend the functionality
of a backend implementation to suit your modeling requirements. This versatility is advantageous in
numerous ways, such as supporting custom operations like flash attention or manually fusing
operations for enhanced performance.

In this section, we will go into the process of extending a backend, providing multiple examples.
But before we proceed, let's establish the fundamental principles that will empower you to craft
your own backend extensions.

Burn's user-facing tensors and modules are runtime-dispatched and don't expose a backend generic.
Backend traits remain part of the lower layer, where they define primitive operations that can
be registered with the Tensor → Bridge → Dispatch → Backend stack. To create an extension, define a
backend trait specifying the new primitive operation, implement it for the backends you support,
and expose a backend-independent Tensor function that calls through Dispatch.

pub trait Backend: burn::backend::Backend {
    fn my_new_function(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
        // You can define a basic implementation reusing the Burn Backend API.
        // This can be useful since all backends will now automatically support
        // your model. But performance can be improved for this new
        // operation by implementing this block in specific backends.
    }
}

You can then implement your new custom backend trait for any backend that you want to support:

impl Backend for burn_tch::LibTorch {
   fn my_new_function(tensor: TchTensor) -> TchTensor {
      // My Tch implementation
   }
}

impl Backend for burn_flex::Flex {
    // No specific implementation, but the backend can still be used.
}

You can support the backward pass using the same pattern.

impl Backend for burn_autodiff::Autodiff {
    // No specific implementation; autodiff will work with the default
    // implementation. Useful if you still want to train your model, but
    // observe performance gains mostly during inference.
}

impl Backend for burn_autodiff::Autodiff {
   fn my_new_function(tensor: AutodiffTensor) -> AutodiffTensor {
      // My own backward implementation, generic over my custom Backend trait.
      //
      // You can add a new method `my_new_function_backward` to your custom backend
      // trait if you want to invoke a custom kernel during the backward pass.
   }
}

impl Backend for burn_autodiff::Autodiff {
   fn my_new_function(tensor: AutodiffTensor) -> AutodiffTensor {
      // My own backward implementation, generic over a backend implementation.
      //
      // This is another way to call a custom kernel for the backward pass that
      // doesn't require the addition of a new `backward` function in the custom backend.
      // This is useful if you don't want all backends to support training, reducing
      // the need for extra code when you know your model will only be trained on one
      // specific backend.
   }
}

The specifics of each implementation will be covered by the examples provided in this section. The
cubecl compiler frontend is the recommended method of implementing custom kernels, since it
supports multiple backends, including wgpu and CUDA, and is the way first-party burn kernels
are written.

Passing structs and enums of tensors Copy Chapter

An extension operation is not limited to individual tensor arguments. A custom struct or enum whose
fields are tensor primitives can be passed to and returned from an operation by deriving
ExtensionType. Fields that are not tensors pass through unchanged, and a field that is itself an
ExtensionType can be nested by annotating it with #[extension_type].

use burn::backend::{
    ExtensionType, backend_extension,
    tensor::{FloatTensor, IntTensor},
};

#[derive(ExtensionType)]
pub struct Boxes {
    pub coords: FloatTensor,
    pub scores: FloatTensor,
    pub count: usize, // Non-tensor fields pass through unchanged.
}

#[derive(ExtensionType)]
pub enum Operand {
    Dense(FloatTensor),
    Sparse { values: FloatTensor, indices: IntTensor },
    Empty,
}

Such a type can be returned from an operation directly. To pass one as an input, mark the argument
with #[extension_type]:

#[backend_extension(Wgpu, Cuda, Autodiff)]
pub trait Backend: burn::backend::Backend {
    // Struct as an output.
    fn detect(image: FloatTensor<Self>) -> Boxes<Self>;

    // Struct or enum as an input.
    fn nms(#[extension_type] boxes: Boxes<Self>, iou_threshold: f32) -> Boxes<Self>;
}

Inputs marked this way can be freely mixed with plain tensor arguments and with each other, and an
operation can take several of them. The backend is selected by looking at a representative tensor
across the inputs, so an enum currently on a variant that holds no tensor simply defers to the next
input; if no input holds a tensor at all, the backend cannot be resolved and the operation panics.

Struct and enum inputs also work with Autodiff. Float fields carry the gradient; other fields do
not. Your impl ... for Autodiff writes the backward pass by hand, exactly as it does for plain
tensor inputs.

File: burn-book/src/advanced/README.md Copy Chapter

Advanced

In this section, we will go into advanced topics that extend beyond basic usage. Given Burn's
exceptional flexibility, a lot of advanced use cases become possible.

Before going through this section, we strongly recommend exploring the
basic workflow section and the
building blocks section. Establishing a solid understanding of how
the framework operates is crucial to comprehending the advanced concepts presented here. While you
have the freedom to explore the advanced sections in any order you prefer, it's important to note
that this section is not intended to be linear, contrary to preceding sections. Instead, it serves
as a repository of use cases that you can refer to for guidance as needed.

File: contributor-book/src/guides/adding-a-new-operation-to-burn.md Copy Chapter

Adding a New Operation to burn

Let's discuss how one might go about adding new operators to Burn, using the example of the pow
operator added in this PR.

Adding the Op to burn-tensor Copy Chapter

burn-tensor is the crate that defines all tensor operations that need to be implemented by the
various backends. The core of this lies in
crates/burn-backend/src/tensor/ops/numeric.rs,
which is home to the numeric trait. The numeric trait is the home of all tensor operations that are
numeric in nature and that are shared by Int and Float Tensor types. The numeric trait is
implemented in
crates/burn-backend/src/tensor/ops/int.rs
for the int type and in
crates/burn-backend/src/tensor/ops/float.rs
for the float type. More information on the relationship between Tensor modules can be found under
the section for Tensor Architecture.

Here is where pow was added to crates/burn-tensor/src/tensor/api/numeric.rs:

1. for the
   Tensor struct
2. for the
   numeric trait
3. for the implementation of numeric for
   float
   and
   int

Tensor is a struct that has a single member: primitive (defined
here),
that is defined by its
Kind:
one of Bool, Float, or Int (those linked in 3). These call the ops for that data type defined
in the
Backend
supertrait[^supertrait]. This is the trait that is then implemented by the different burn-
backends (such as burn-flex and burn-wgpu) which must implement the functions if no default
is provided.

In this case, we don't need to worry about Bool Tensors. Float ops are implemented under
crates/burn-backend/src/backend/ops/tensor.rs,
and Int ops under
crates/burn-backend/src/backend/ops/int_tensor.rs.
The current convention is ops of each type, if not unique to that type, are prefixed with the type.
So powf and sundry would be defined as int_powf for IntTensorOps and float_powf for
FloatTensorOps. If an op is unique to a type, then it should be implemented under
burn-tensor/src/api/{type}.rs. For example, here is an implementation for
sin under crates/burn-tensor/src/api/float.rs
which obviously doesn't make sense for Int or Bool tensors.

The Int Tensor function uses the ones defined for Float with 2 extra casts (LHS to a Float
tensor, Output to an Int). Given that the rest of the code will only look at the float
implementations.

With the addition of quantized float tensors, the Float tensor primitive is represented by the
TensorPrimitive
enum. This allows us to handle both float and quantized float operations in the Tensor
implementation, correctly dispatching to the corresponding op (float or quantized) based on the
variant. Following the same convention, the equivalent
quantized tensor ops
are prefixed with q_* (e.g., q_reshape instead of float_reshape). Most ops have a default
implementation that simply dequantizes the input into its floating-point representation, performs
the operation on the float tensor, and quantizes the output. Backends can overwrite specific
implementations when required/desired.

Adding Tests Copy Section

Additional tests should be added to burn-backend-tests under
crates/burn-backend-tests/tests/tensor/{float_or_int}/ops/{op_name}.rs,
and the module name should be inserted into
crates/burn-backend-tests/tests/tensor/{float_or_int}/ops/mod.rs.

If it makes sense for a floating point operation to support quantization, the
QTensorOps
counterpart is usually added at the same time with a default implementation (as mentioned in the
previous section). Tests for q_* ops follow a similar procedure: the test is added under
crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/{op_name}.rs,
the module name is inserted into
crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/mod.rs.
If you take a look at any of the existing tests for an operation on a quantized tensor, you will see
that the inputs and expected outputs are always defined with floating point values. While it assumes
that the quantization and dequantization are correct, it makes the tests much more readable and
easier to understand w.r.t. what is being tested. Effectively, the tests are there to ensure that a
tensor operation is invariant to quantization (up to some quantization error, of course).

Note: the tests try to use tensors with floating point values which can be de/quantized without
introducing too much quantization error, but the result always depends on the operation (e.g.,
tensor product of values can grow larger and significantly increase the output tensor range, leading
to more de/quantization error on the results).

Adding the Op to burn-autodiff Copy Chapter

Since this is probably the hardest and the least straightforward, we'll cover this backend
separately. burn-autodiff enables other backends to use autodifferentiation[^autodiff]. Ops for
float types are implemented in
crates/burn-autodiff/src/ops/tensor.rs
and need to:

1. Define a unit struct [^absolute_units] that implements a backward (pass) function
2. Within the backward function, as this is an elementwise binary operation it implements the binary
   function (from backward.rs under the same directory), the last 2 arguments are two closures
   that define the left and right partial derivatives.
3. Then define what happens when a specific operation is tracked or untracked, where untracked just
   calls the function in the normal way, and tracked sets the execution the backward function
   defined above.
4. When tracked, operations are part of the autodiff graph and must save the needed information to
   efficiently perform their backward pass later. If the information is light (such as a shape), it
   should be directly saved in the state. If the operation's inputs are needed to compute the
   backward pass, it should be checkpointed rather than saved. This will allow the input to be
   provided lazily at the backward pass depending on the checkpointing strategy.
5. An operation must also be identified as compute-bound (.computeBound()) or memory-bound
   (.memoryBound()) for gradient checkpointing. Compute-bound operation are heavy to compute
   (for instance matmul or convolution), which means that even with checkpointing they will save
   their output for the backward pass and not recompute it. Memory-bound operations are more
   trivial (like powf which only performs one small operation per tensor entry), so it can be
   beneficial to recompute them during the backward pass instead of saving their whole forward
   output to memory. Operations registered as memory-bound need to know their parents
   (.parents() method) and how to recompute their forward pass during the backward pass (with a
   struct that implements RetroForward), using their parents' outputs.

The above steps are mostly boilerplate, so you can often just copy the contents of another similar
op, change the name of the structs, and ensure that either both sides have the data they need (if
they need to have a copy of the opposite sided tensor, clone its contents).

Computing derivatives Copy Section

For those that need it, here is a quick refresher on the necessary calculus. If you are familiar
with how to calculate partial derivatives, you can skip this section.

Since pow is a binary operation, the left and right functions are the partial derivatives with
respect to the left and right sided tensors.

Let's define the operator as a function \(f(x,y)=x^{y}\) , where \(x\) is the left hand tensor
and \(y\) is the right handed tensor. The two closures are defining the partial derivatives of
\(f\) with respect to \(x\),\(y\). Treat the other variables as a constant

$$\frac{\delta }{\delta x} (x^{y})= y \cdot x^{y-1}$$ is the left handed closure, and

$$\frac{\delta }{\delta y} (x^{y}) = x^{y} \cdot ln(x)$$

is the right. If you aren't sure how to calculate these by hand, it is recommended to use
symbolab,
plug in your operator in terms of \(x\) and \(y\), and just swap out the variable
\(x\)|\(y\) in the partial derivative to get the other side.

Testing autodiff Copy Section

For testing the autodiff operations, please refer to
this section.

Adding the Op to other backends Copy Chapter

Most of these are fairly straightforward implementations. For reference here's pow's float
implementation for torch and flex backends:

1. Torch implementation in
   crates/burn-tch/src/ops/tensor.rs
   and the Op used in
   crates/burn-tch/src/ops/base.rs
2. Flex in
   crates/burn-flex/src/ops/float.rs

This is where any calculation happens currently. Playing a guessing game with method names and
seeing what completions are suggested will take you far. If you are having trouble figuring out how
to do it from the docs for that backend,
try searching github for relevant function calls.

Adding the Op to fusion, JIT and cubecl backends Copy Chapter

Adding an operator to these backends can be fairly straightforward, though due to what these
backends are for, involves a bit more indirection. Fusion and jit, like autodiff, are not target
backends as much as backends that enable certain functionality for other backends, in this case
kernel fusion or just-in-time compilation. Adding the operator won't involve doing any calculation,
you'll just be describing how the generated code should look. Most of this can be
copy/pasted/adjusted from other functions.

Here's how powf was added to burn-fusion:

1. Added powf to the float ops under
   crates/burn-fusion/src/ops/tensor.rs
2. Added powf to the NumericOperationIr enum under
   crates/burn-ir/src/operation.rs
3. Added powf to the implementations of NumericOperationIr enum under
   crates/burn-ir/src/operation.rs
4. Added powf to the implemented of NumericOperationIr enum under
   burn/crates/burn-fusion/src/stream/context.rs

The way cubecl handles tensor-scalar operations is by transforming both into a sequence of
vectorized scalar operations. Since powf already existed in cubecl, it was pretty easy to reuse
the existing implementation for the situation where both sides of the operation were tensors. The
cubecl crate is primarily concerned with how the operation is compiled and executed by the gpu.
The actual implementation is defined in burn-cubecl.

Here is where code was added for powf in burn-cubecl and cubecl:

1. to the implementation of
   FloatTensorOps under burn/crates/burn-cubecl/src/ops/tensor.rs
2. the function being called was added to
   burn/crates/burn-cubecl/src/ops/numeric.rs
3. the operator was defined in
   cubecl/crates/cubecl-ir/src/arithmetic.rs
4. how the operation looks to the gpu was added to
   burn/crates/burn-cubecl-fusion/src/engine/codegen/ir.rs
5. the mappings between the gpu operation and the CPP, WGSL and SPIR-V instructions were added to
   cubecl/crates/cubecl-cpp/src/shared/base.rs,
   cubecl/crates/cubecl-wgpu/src/compiler/wgsl/compiler.rs
   and
   cubecl/crates/cubecl-spirv/src/arithmetic.rs
6. the instructions themselves were added for WGSL to
   instruction op enum in cubecl/crates/cubecl-wgpu/src/compiler/wgsl/instructions.rs,
   and the actual
   instruction in wgsl here,
   for CPP in the enum here
   cubecl/crates/cubecl-cpp/src/shared/instruction.rs
   and the actual instruction here
   cubecl/crates/cubecl-cpp/src/shared/binary.rs

We needed to generate some custom WGSL code for powf in WGSL, primarily due to issues with proper
case handling of the wgsl pow function, like 0 to the 0 power being 1, and any negative number to an
even power being positive. We reused as much as the existing logic as possible, and then branched at
the last point based off the var type of the rhs.
See here.
For most operations, you shouldn't need to add to cubecl-wgpu/src/compiler/wgsl/extension.rs
unless the operation isn't native to WGSL.

For functions that need a complex kernel without a direct mapping to a base instruction, simply use
the cube macro (see
the cubecl book).

And you're done! Congrats, you just fully added a new operation to burn, and we are all one step
closer to the answer to Are we learning yet? being "Yes, and
it's freaking fast!". Buy yourself a coffee.

[^supertrait]:
for more on supertraits see
the advanced trait section of the rust book

[^autodiff]:
wiki link for
automatic differentiation

[^absolute_units]:
for more information on unit structs see
the defining and instantiating structs section of the rust book

File: contributor-book/src/guides/README.md Copy Chapter

Guides for Contributors

The following guides are meant to help contributors accomplish specific tasks, such as adding new operations to Burn.

File: contributor-book/src/guides/submitting-examples.md Copy Chapter

Submitting Examples to Burn

This guide explains how to create and submit new examples to the Burn repository. Examples are a great way to demonstrate Burn's capabilities and help users understand how to use the framework effectively.

For a minimal working example, see the simple-regression example in the repository.

Repository Structure Copy Chapter

The Burn repository is set up as a workspace, with examples located in the examples/ directory. Each example is a separate crate that can reuse workspace dependencies.

Creating a New Example Copy Chapter

1. Navigate to the examples directory:

   bash Copy

   cd examples

2. Create a new library crate:

   bash Copy

   cargo new --lib <my-example>

3. Update the example's Cargo.toml:

   toml Copy

   [package]
   name = "<my-example>"
   version = "0.1.0"
   edition = "2021"
   readme = "README.md"
   # Remove this line if it exists
   # readme.workspace = true
   
   [dependencies]
   # Reuse workspace dependencies when available
   serde = { workspace = true }
   # Add example-specific dependencies
   burn = { path = "../../" }

Required Files and Structure Copy Chapter

Each example must include a README.md file with:

• A brief description of what the example demonstrates
• A terminal command showing how to run the example
• Any prerequisites or setup instructions

Example README structure:

# Example Name

Brief description of what this example demonstrates.

## Running the Example

```bash
cargo run --example <my-example>
```

## Prerequisites

List any prerequisites here.

Source Code Structure Copy Section

• src/ directory: Contains the main implementation code
• examples/ directory: Contains example code
  • <my-example>.rs: Example implementation

Resource Handling Copy Chapter

• Resources (datasets, models, etc.) should be downloaded in the example code
• Do not track external files in the repository
• Include code to download and prepare resources when the example is run

Best Practices Copy Chapter

1. Code Organization

   • Keep the code modular and well-documented
   • Use clear, descriptive variable and function names
   • Include comments explaining complex operations

2. Error Handling

   • Implement proper error handling
   • Provide meaningful error messages
   • Handle resource download failures gracefully

3. Performance

   • Optimize for reasonable execution time
   • Include progress indicators for long-running operations
   • Consider adding configuration options for different hardware capabilities

4. Documentation

   • Document all public APIs
   • Include inline comments for complex logic
   • Explain any non-obvious implementation details

Submitting Your Example Copy Chapter

1. Ensure your example follows all the guidelines above
2. Test your example thoroughly
3. Create a pull request with:
   • A clear description of what the example demonstrates
   • Any relevant issue numbers
   • Screenshots or output examples (if applicable)

Feel free to ask questions in the pull request if you need clarification or guidance.