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 a property of the device rather than a separate type parameter. Enable the autodiff feature to add automatic differentiation support to devices.

Backend selection

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.

Autodiff

Requires autodiff feature.

Gradient computation is opt-in for a device:

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

// Tensors created on this device will track gradients
let x = Tensor::<1>::from_floats([1.0, 2.0, 3.0], &device);

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 inner(self) -> Self

Returns the underlying device, removing the autodiff capability if present.

If autodiff is not enabled, this method returns the device as-is.

Example

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

assert!(!inner_device.is_autodiff());

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)

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.

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 true if autodiff (gradient tracking) is enabled on this device.

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 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 and integer data types used when creating tensors on this device.

See configure to configure them.

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, and must happen before any tensor is created on the device. The first tensor operation will lock the device to its defaults, causing subsequent initializations attempt to return DeviceError::AlreadyInitialized.

Errors

Returns DeviceError::AlreadyInitialized if settings have already been set for this device (either by a prior call or because a tensor operation has already occurred).

Example

let device = Default::default();

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.

Local backends (CPU, CUDA, WGPU, …) enumerate the hardware found on the host. The Remote variant instead lists every device hosted by the burn-remote server at the given address — it connects to the server to learn how many devices it exposes:

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

// Every device hosted by a remote server.
let remote = Device::enumerate(DeviceType::remote_websocket("ws://host:3000"));

// Filters combine with `|`.
let both = Device::enumerate(DeviceType::Cuda | DeviceType::remote_websocket("ws://host:3000"));

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 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. To check if two devices have identical capabilities, check Device::is_autodiff.

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.