Struct Device 

pub struct Device { /* private fields */ }

A high-level device handle for tensor operations.

Device provides a unified interface to interact with the underlying compute backend.

Autodiff support is configured on the device rather than through a separate type parameter. Tensors inherit that context when created and can later change it independently. Enable the autodiff feature to add automatic differentiation support to devices.

Backend selection

Default Cargo features do not enable an execution backend. Select one explicitly, for example with the wgpu or flex feature; there is no implicit CPU fallback. Backend-free builds can expose tensor/model APIs, but cannot create an execution device.

Device::default() selects the first enabled backend in this order: CUDA, Metal, ROCm, Vulkan, WebGPU, wgpu, CPU, LibTorch, Flex, Remote, NdArray. Use an explicit factory method when the choice must be independent of Cargo feature unification. Without an execution backend, Device::default() panics with configuration guidance.

Enable the desired backend via Cargo feature flags, then call the corresponding factory method:

// Default CUDA device (requires the `cuda` feature).
let device = Device::cuda(DeviceIndex::Default);

// CUDA device at hardware index 1.
let device = Device::cuda(1);

// WGPU with explicit selector (requires `wgpu`/`vulkan`/`metal`/`webgpu`).
let device = Device::wgpu(DeviceKind::DiscreteGpu(0));

// Default device for whichever backend is enabled.
let device = Default::default();

Available factory methods (each gated by its matching Cargo feature): Device::cpu, Device::cuda / Device::rocm / Device::libtorch_cuda (take an integer index or a DeviceIndex), Device::wgpu / Device::vulkan / Device::metal / Device::webgpu (take a DeviceKind), Device::flex, Device::ndarray, Device::libtorch, Device::libtorch_mps, Device::libtorch_vulkan, Device::capture.

Autodiff

Requires autodiff feature.

Gradient computation is opt-in for a device:

let device = Device::default().autodiff();

// Tensors created on this device can participate in an autodiff graph.
let x = Tensor::<1>::from_floats([1.0, 2.0, 3.0], &device).require_grad();

Implementations

impl Device

pub fn new(device: impl Into<DispatchDevice>) -> Self

Wrap a backend-specific device in a unified Device.

Used by:

• the backend-specific factory methods below (Device::cuda, etc.) — these are the recommended entry points for downstream code;
• burn-tensor’s bridge ops, which already hold a DispatchDevice and just need to wrap it;
• direct callers (tests, type-erased helpers) that have a concrete backend device type at hand.

Anything convertible into DispatchDevice is accepted, including DispatchDevice itself.

pub fn as_dispatch(&self) -> &DispatchDevice

Borrow the underlying DispatchDevice.

The inverse of Device::new. Useful to backend-extension authors who need to dispatch on the concrete backend variant (e.g. matching DispatchDevice::Remote(_)).

impl Device

pub fn cpu() -> Self

Available on crate feature cpu only.

Default CPU device backed by CubeCL’s CPU backend.

pub fn cuda(index: impl Into<DeviceIndex>) -> Self

Available on crate feature cuda only.

CUDA device at the given hardware index.

Accepts a plain integer (e.g. Device::cuda(0)) or a DeviceIndex — use DeviceIndex::Default to let the backend pick.

pub fn rocm(index: impl Into<DeviceIndex>) -> Self

Available on crate feature rocm only.

ROCm/HIP device at the given hardware index.

Same selector semantics as Device::cuda.

pub fn flex() -> Self

Available on crate feature flex only.

Flex backend device.

pub fn wgpu(device_kind: DeviceKind) -> Self

Available on crate feature wgpu only.

WGPU device, selected via DeviceKind.

AutoCompiler dispatches to the most appropriate shader language (WGSL, SPIR-V, or MSL) based on the enabled features.

For DeviceKind::DefaultDevice, the adapter is picked by wgpu’s selection heuristics (high-power GPU preferred, or overridden by CUBECL_WGPU_DEFAULT_DEVICE).

The graphics API — and so the shader compiler — is the runtime’s to settle, from the enabled features and what the machine offers. Device::vulkan, Device::metal and Device::webgpu pin one instead: the same adapter on two APIs is two devices.

This constructor is lazy. Use Device::wgpu_options to initialize with runtime options or share an existing wgpu setup. In a browser, initialize asynchronously with init_async before creating tensors.

pub fn vulkan(device_kind: DeviceKind) -> Self

Available on crate feature vulkan only.

Vulkan-backed WGPU device, selected via DeviceKind — and so SPIR-V, where the build and the adapter allow it.

Pinned to Vulkan: the device comes up there or not at all, rather than quietly on whichever API Device::wgpu would settle on.

pub fn metal(device_kind: DeviceKind) -> Self

Available on crate feature metal only.

Metal-backed WGPU device, selected via DeviceKind, requiring native MSL support.

Pinned to Metal, as Device::vulkan is to Vulkan. Runtime initialization panics if native MSL is unavailable on the selected device. Use Device::wgpu for automatic compiler selection with WGSL fallback.

pub fn webgpu(device_kind: DeviceKind) -> Self

Available on crate feature webgpu only.

WebGPU-backed device, selected via DeviceKind — the browser’s own. Pinned to WebGPU, as Device::vulkan is to Vulkan.

pub fn without_autodiff(self) -> Self

Returns this device without its autodiff association.

If autodiff is not enabled, the device is returned unchanged. This operation is idempotent.

Example

let device = Device::default().autodiff();
let inference_device = device.without_autodiff();

assert!(!inference_device.is_autodiff());

pub fn inner(self) -> Self

Returns this device without its autodiff association.

This is equivalent to without_autodiff. inner reflects the historical backend-decorator model, while without_autodiff describes the device’s runtime property directly.

pub fn sync(&self) -> Result<(), ExecutionError>

Synchronize the device, waiting for all pending operations to complete.

Errors

Returns an ExecutionError if an operation failed to execute.

pub fn flush(&self) -> Result<(), ExecutionError>

Flush the device’s pending operations, handing them off for execution without waiting for them to complete.

Backends that buffer work hold registered operations in a local queue until enough accumulate: the fusion backend batches ops to build optimizations, and the remote backend batches them before sending them over the network. flush forces that queue out now — the fusion backend processes its pending optimizations and the remote backend sends its batch to the server.

Unlike sync, this does not block on results — it only ensures buffered operations are dispatched instead of sitting idle. Eager backends, which execute each operation as it is registered, have nothing buffered and treat this as a no-op.

Errors

Returns an ExecutionError when the buffered operations cannot be dispatched, e.g. on a device that is poisoned.

pub fn profile<O: Send + 'static>( &self, func: impl FnOnce() -> O + Send, ) -> Result<(O, ProfileDuration), ExecutionError>

Measure how long this device spends on the work func puts on it, in device time.

The measurement is a ProfileDuration: a future the device answers once it has run both ends of the window, so nothing here waits on the device, and windows nest — an inner profile costs the outer one nothing. Collect them and resolve once the run is over.

The window spans the stream from the call to func’s return. Work the stream still owed from before falls in; work a backend queues past the end falls out — the fusion backend holds a closure’s last operations back to batch them, so a window over lazy work alone can read as empty. Ending the closure with a read, or profile_with and ProfileOptions::flush, closes the window over all of it. A window that nothing ran in reads as no time.

A backend with no device clock (ndarray, LibTorch, a remote device whose server has none) measures wall-clock time between two syncs instead: that one waits, and an inner window’s syncs are charged to the outer.

let (output, duration) = device.profile(|| model.forward(input))?;
// Later, once the run is over:
let ticks = duration.resolve().await.expect("the window carried work");
println!("forward: {:?}", ticks.duration());

Errors

Returns an ExecutionError when the device refuses to open or close the window — a remote device does from a browser thread, which cannot wait on the server.

pub fn profile_with<O: Send + 'static>( &self, options: ProfileOptions, func: impl FnOnce() -> O + Send, ) -> Result<(O, ProfileDuration), ExecutionError>

profile with ProfileOptions.

pub fn seed(&self, seed: u64)

Seeds the random number generator for this device.

Seeding before tensor operations that involve randomness (e.g. Tensor::random) makes those operations reproducible in a single-threaded program.

Note

Depending on the backend, the seed may be applied globally rather than scoped to this specific device. It is guaranteed that at least this device will be seeded.

Example

let device = Default::default();
device.seed(42);
let t = Tensor::<1>::random([8], Distribution::Default, &device);

pub fn is_autodiff(&self) -> bool

Returns whether this device is associated with autodiff.

This is device context inherited by newly created tensors, not a statement that any tensor participates in a graph or retains gradients.

Example

let device = Default::default();
assert!(!device.is_autodiff());

let ad_device = device.autodiff();
assert!(ad_device.is_autodiff());

pub fn supports_dtype(&self, dtype: impl Into<DType>) -> bool

Returns true if this device supports dtype for general computation: storage, conversion, and arithmetic.

A type can be less than generally supported — bf16 on a Vulkan device, for example, is often storable and convertible but has no arithmetic (SPIR-V’s SPV_KHR_bfloat16 permits only conversions, dot products, and cooperative-matrix use). Computing in such a type produces backend-dependent garbage, so check before selecting a reduced precision:

let dtype = if device.supports_dtype(FloatDType::BF16) {
    FloatDType::BF16
} else {
    FloatDType::F32
};

pub fn memory_persistent_allocations<Output: Send, Input: Send, Func: Fn(Input) -> Output + Send>( &self, input: Input, func: Func, ) -> Output

Sets the current allocation mode to persistent.

pub fn memory_cleanup(&self)

Triggers a memory cleanup on this device.

The amount of memory reclaimed depends on the allocator implementation. Calling this method does not guarantee that any memory will be freed.

pub fn memory_pool_report(&self) -> Option<Vec<SlicedPoolReport>>

This device’s dynamic pools, in the order they were installed. None on a backend that does not report them.

pub fn memory_pool_usage(&self) -> Option<MemoryPoolUsage>

What this device’s allocator currently holds. None on a backend that does not report it.

pub fn staging<'a, Iter>(&self, data: Iter)

where Iter: Iterator<Item = &'a mut TensorData>,

Prepares the given data for transfer between the CPU and accelerator devices such as GPUs.

Depending on the backend, the data may be transferred to pinned memory or another transfer-optimized format to improve transfer performance.

pub fn settings(&self) -> DeviceSettings

Returns the DeviceSettings for this device.

Settings include the default float, integer, and boolean data types used when creating tensors on this device.

Before initialization, returns a snapshot of the backend defaults without locking them. Another thread may configure the device afterward, so subsequent tensor operations may use different settings. Configure the device before relying on its defaults.

pub fn configure( &mut self, config: impl Into<DeviceConfig>, ) -> Result<(), DeviceError>

Configures the settings for this device.

This configures the dtype used when no explicit type is specified at tensor creation time.

Settings can only be initialized once per device. Configure defaults before creating tensors or initializing model parameters: tensor creation locks them to the backend’s defaults, even with an explicit dtype, and any later call returns DeviceError::AlreadyInitialized. Querying settings does not lock them.

Individual tensors can still use an explicit supported dtype at creation or be converted with Tensor::cast; neither changes the defaults.

Errors

Returns DeviceError::UnsupportedDType if a requested dtype is unsupported. Returns DeviceError::AlreadyInitialized if settings have already been initialized for this device, either by a prior call or by a tensor operation.

Example

use burn_tensor::{Device, FloatDType, Int, IntDType, Tensor};

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

// Float tensors will now use F16
let floats = Tensor::<2>::zeros([2, 3], &device);
// Int tensors will now use I32
let ints = Tensor::<2, Int>::zeros([2, 3], &device);

pub fn enumerate(filter: impl Into<DeviceFilter>) -> Devices

Retrieves all available Devices that match the given DeviceType filter.

Backends enumerate the hardware found on this machine, and DeviceType::Remote every device a remote server hosts:

// Every CUDA device on this machine.
let local = Device::enumerate(DeviceType::Cuda);

// Filters combine with `|`.
let both = Device::enumerate(DeviceType::Cuda | DeviceType::Remote(host));

Panics

Where RemoteHost::devices returns an error for a DeviceType::Remote, and on wasm for any DeviceType::Remote, which a browser can only list with RemoteHost::devices_async.

pub fn identity(&self) -> Option<DeviceIdentity>

Available on crate feature cubecl only.

Who this device is: the part’s name, what its kernels are compiled for and, for a card on a bus, the card itself. None for a backend that does not report one. Opens the device.

pub fn enumerate_physical() -> Vec<PhysicalGpu>

Available on crate features cpu or cuda or rocm or wgpu only.

Every card this build can reach, once each, with the devices that reach it.

A card visible to two runtimes (an NVIDIA GPU under CUDA and under Vulkan) is one entry holding both devices, so a caller placing work on cards never puts two stages on one card by another name. Opens every device of every runtime, so call it once and keep the answer.

pub fn performance_stats(&self, keys: &[ThroughputKey]) -> Vec<ThroughputStat>

Available on crate feature cubecl only.

Measure peak compute and memory throughput for this device.

Runs cubecl-std’s throughput benchmarks for each ThroughputKey, returning one ThroughputStat per key (in the same order). A key the device has no peak for keeps its place, carrying the ThroughputError saying why. Only cubecl-backed devices (cuda, wgpu, …) report measurements; other backends return an empty vector.

impl Device

pub fn wgpu_options() -> WgpuOptions

Available on crate feature wgpu only.

Options for initializing Burn’s wgpu runtime or registering an existing setup.

Starts with the default device selector, automatic graphics API, and runtime defaults. Call init or init_async before creating tensors on the same selector. An initialized runtime cannot be reconfigured. Clone its device to reuse it.

Trait Implementations

impl Clone for Device

fn clone(&self) -> Self

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for Device

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl Default for Device

fn default() -> Self

Returns the “default value” for a type.

impl Eq for Device

impl From<&Device> for TensorCreationOptions

fn from(device: &Device) -> Self

Convenience conversion from a reference to a device.

Example:

use burn_tensor::TensorCreationOptions;
use burn_tensor::Device;

fn example(device: Device) {
    let options: TensorCreationOptions = (&device).into();
}

impl<D: Into<DispatchDevice>> From<D> for Device

fn from(device: D) -> Self

Converts to this type from the input type.

impl FromIterator<Device> for Devices

fn from_iter<I: IntoIterator<Item = Device>>(devices: I) -> Self

Creates a value from an iterator.

impl PartialEq for Device

fn eq(&self, other: &Self) -> bool

Compares devices based on hardware identity.

Returns true if both devices represent the same compute resource. Note that this comparison ignores autodiff and checkpointing settings. Inspect Device::is_autodiff and Device::gradient_checkpointing_strategy() when execution context also matters.

fn ne(&self, other: &Self) -> bool

Compares devices based on hardware identity.

Returns false if both devices represent the same compute resource, even if one has autodiff enabled and the other does not.