Crate autodiff 

Burn Autodiff

Reverse-mode automatic differentiation as a backend decorator.

Autodiff wraps any backend B and records the operations needed to compute gradients. Only first-order derivatives are supported.

Most applications do not name this type. Enable the autodiff feature of burn (also enabled by train) and turn autodiff on for a device:

let device = Device::wgpu(Default::default()).autodiff();
let x = Tensor::<2>::ones([2, 2], &device).require_grad();
let grads = (x.clone() * 3.0).sum().backward();
let x_grad = x.grad(&grads).unwrap();

Dispatch then routes operations on that device through Autodiff. Use this crate directly when implementing a backend extension that needs custom backward passes.

Gradient checkpointing

The second type parameter selects a CheckpointStrategy. NoCheckpointing keeps every activation needed by the backward pass. BalancedCheckpointing recomputes cheap operations instead of storing their outputs, trading compute for memory. At the Device level, device.autodiff().gradient_checkpointing() selects it.

Feature flags

• std (default): standard library support. Without it the crate is no_std with alloc.
• tracing: instrument operations with the tracing crate.

Modules

checkpoint

Checkpoint module.

distributed

Distributed utils.

grads

Gradients module.

ops

Operation module.

Macros

retro_binary

Creates a RetroForward struct for binary operations

retro_unary

Creates a RetroForward struct for unary scalar operations

retro_unary_scalar

Creates a RetroForward struct for unary scalar operations

Structs

Autodiff

Enable auto-differentiation on a backend.

NodeId

Unique identifier generated for each node.