Backend and Device

Burn's user-facing API is centered on Tensor, Module, and Device. These types are not generic over a backend. Instead, every tensor carries a runtime Device that identifies where and how its operations execute.

use burn::tensor::{Device, Tensor};

let device = Device::wgpu(Default::default());
let tensor = Tensor::<2>::ones([2, 3], &device);

// The same model and tensor API can target another backend.
let other_device = Device::cuda(0);
let tensor = tensor.to_device(&other_device);

The corresponding Cargo feature must be enabled for each device constructor. For example, the wgpu and cuda features make Device::wgpu and Device::cuda available.

Selecting a Device

Device provides constructors for the backends enabled in your build. Multiple backend features can be enabled together; the device chooses where operations execute.

The vulkan, metal, and webgpu features also enable wgpu. Enabling a feature makes that option available; it does not force Device::wgpu to use it. Automatic devices retain WGSL fallback even with native compiler features enabled. The explicit constructors pin the graphics API, so Device::vulkan cannot fall back to a different API even when it uses WGSL instead of SPIR-V.

For example, one build can support both automatic selection and explicit Metal selection:

burn = { version = "0.22", features = ["vulkan", "metal"] }

use burn::tensor::Device;

let automatic = Device::wgpu(Default::default());
let metal = Device::metal(Default::default());

The explicit Metal device requires an available Metal adapter with native MSL support. The automatic device selects an available graphics API and compiler for the platform.

Indexed devices accept either an integer or DeviceIndex. WGPU-family constructors accept a DeviceKind, which can select a discrete, integrated, or virtual GPU, a CPU adapter, or the best available device:

use burn::tensor::{Device, DeviceIndex, DeviceKind};

let default_wgpu = Device::wgpu(Default::default());
let discrete_wgpu = Device::wgpu(DeviceKind::DiscreteGpu(0));
let integrated_wgpu = Device::wgpu(DeviceKind::IntegratedGpu(0));
let default_cuda = Device::cuda(DeviceIndex::Default);
let second_cuda = Device::cuda(1);

Burn also supports remote devices, hosted by a Burn server in another process on this machine or another, with the remote feature. Device::remote_options(&host).init() connects one, where a RemoteHost names the server: see Distributed Computing.

Using a Device

Tensor creation methods take &Device. Modules receive the same device when their parameters are initialized:

let device = Device::cuda(0);
let input = Tensor::<2>::zeros([32, 128], &device);
let model = ModelConfig::new().init(&device);
let output = model.forward(input);

Use Tensor::to_device or Module::to_device to move existing values. Operations involving multiple tensors require compatible devices, so move them explicitly before combining them.

let cpu = Device::flex();
let gpu = Device::cuda(0);

let tensor = Tensor::<2>::ones([2, 3], &cpu);
let tensor = tensor.to_device(&gpu);
let model = model.to_device(&gpu);

To move a trainable module and optimize its parameters on the destination, use Module::fork. Module::to_device preserves connections to the source parameters instead. See the module transfer semantics.

Autodiff and Execution

Devices provide the autodiff and checkpointing defaults for newly created tensors. Each tensor carries its own context and can later change it independently. Tensor::to_device preserves the source tensor's context, regardless of the destination's defaults. Calling autodiff returns a device whose newly created tensors have autodiff enabled:

let device = Device::wgpu(Default::default());
let training_device = device.autodiff();

assert!(training_device.is_autodiff());
let inference_device = training_device.without_autodiff();
assert!(!inference_device.is_autodiff());

autodiff() and without_autodiff() are idempotent. The historical inner() method is equivalent to without_autodiff(). Chain autodiff().gradient_checkpointing() to enable autodiff with the balanced checkpointing strategy. Repeating autodiff() doesn't enable higher-order differentiation; Burn currently supports first-order autodiff only.

Device equality compares compute resources and ignores autodiff and checkpointing settings. Use tensor.is_autodiff() to inspect a tensor's association; equality with an autodiff device does not establish that the tensor has autodiff enabled.

The following methods are also useful when coordinating execution:

• seed(seed) seeds random operations on the device.
• sync() waits for queued work and reports an execution error if one occurred.
• flush() submits queued work without waiting for completion.
• is_autodiff() reports whether autodiff is associated with the device.
• gradient_checkpointing_strategy() returns the active strategy, or None without autodiff.
• supports_dtype(dtype) reports whether the device supports a dtype.
• memory_cleanup() asks the backend to release unused cached allocations.

Device Settings

Each device has runtime settings, including its default float, integer, and boolean dtypes. Set them with configure() before creating tensors or initializing model parameters, then inspect them with settings():

use burn::tensor::{Device, DeviceConfig, FloatDType, IntDType};

let mut device = Device::cuda(0);
device.configure(
    DeviceConfig::default()
        .float_dtype(FloatDType::F16)
        .int_dtype(IntDType::I32),
)?;

let settings = device.settings();

Configure defaults before the first tensor operation on that device. Device settings are initialized once; tensor creation locks them to the backend's defaults, even with an explicit dtype, and a later configure() call returns DeviceError::AlreadyInitialized. Querying settings() does not lock defaults: before initialization, it returns a snapshot of the backend defaults, which subsequent configuration may change.

These settings choose defaults, not a single precision for every tensor on the device. You can still specify a supported float or integer dtype at creation or cast an existing tensor:

use burn::tensor::{DType, FloatDType, Tensor};

let tensor = Tensor::<2>::zeros([2, 3], (&device, DType::F32));
let tensor = tensor.cast(FloatDType::F16);

Neither choice changes the device's defaults.

Enumerating Devices

Device::enumerate discovers enabled devices matching a filter. This is useful for selecting hardware dynamically or setting up multi-device training:

use burn::tensor::{Device, DeviceFilter, DeviceType};

let devices = Device::enumerate(
    DeviceFilter::new()
        .with(DeviceType::Cuda)
        .with(DeviceType::Wgpu),
);
let devices = devices.into_vec();

The exact DeviceType variants available depend on the backend features enabled for the application.

A GPU can be reachable through more than one runtime: an NVIDIA card enumerates under DeviceType::Cuda and again under DeviceType::Vulkan. Device::enumerate_physical lists each card once, with every device that runs on it, and Device::identity says which card a device is on:

for gpu in Device::enumerate_physical() {
    let address = gpu.physical.pci_address.map(|address| address.to_string());
    println!("{} at {address:?}: {:?}", gpu.name, gpu.devices);
}

A card is recognized by its PCI address, which CUDA, ROCm and Vulkan report, or by its Windows LUID. A card that reports neither, such as an Apple GPU, is listed as a card of its own. The LUID changes when the machine restarts, so anything stored to recognize a card later should use the PCI address. Both calls open the devices they report on, so enumerate once and keep the answer.

Execution Stack

Under the hood, an operation flows through the Tensor → Bridge → Dispatch → Backend stack:

• Tensor provides the stable, backend-independent API used by applications.
• A bridge converts tensor handles into dispatch values.
• Dispatch selects the implementation associated with the tensor's device.
• The backend executes the primitive operation.

This replaces the backend parameter in the former Tensor<B, D, K> API. Cargo features make backends available, and devices select them at runtime, so one application can use several backends with the same tensor and model types. The opaque bridge keeps their concrete implementation types out of the public tensor representation.

The Backend and AutodiffBackend traits still define the low-level implementation contract, but ordinary application code does not need bounds such as B: Backend. You will mainly encounter those traits when implementing or extending a backend; see Backend Extension.

For the responsibilities of each layer, the reason for the opaque bridge, and an operation walkthrough, see the contributor book's Tensor architecture chapter.