Enum DispatchDevice 

pub enum DispatchDevice {
    Flex(FlexDevice),
    Autodiff(AutodiffDevice),
}

Represents a device for the Dispatch.

Each variant corresponds to a backend that the Dispatch can dispatch operations to.

Example

use burn::DispatchDevice;

// One variant covers every cubecl runtime; the device inside says which.
#[cfg(feature = "cuda")]
let cuda_device = DispatchDevice::Cube(cubecl::Device::Cuda(Default::default()));

#[cfg(feature = "ndarray")]
let ndarray_device = DispatchDevice::NdArray(Default::default());

Variants

Flex(FlexDevice)

The Flex backend device (CPU-only).

Autodiff(AutodiffDevice)

The autodiff enabled backend device.

Implementations

impl DispatchDevice

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

Creates a new DispatchDevice with automatic differentiation enabled.

pub fn autodiff_with_gradient_checkpointing( device: impl Into<DispatchDevice>, checkpointing: GradientCheckpointingStrategy, ) -> DispatchDevice

Creates a new DispatchDevice with automatic differentiation and the provided gradient checkpointing strategy enabled.

pub fn inner(self) -> DispatchDevice

Returns the inner device, without autodiff (when enabled).

Trait Implementations

impl Clone for DispatchDevice

fn clone(&self) -> DispatchDevice

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl Debug for DispatchDevice

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

Formats the value using the given formatter.

impl Default for DispatchDevice

fn default() -> DispatchDevice

Select an enabled backend in this order: CUDA, Metal, ROCm, Vulkan, WebGPU, wgpu, CPU, LibTorch, Flex, Remote, NdArray. BURN_DEVICE overrides this in std builds. Capture devices must be constructed explicitly.

The metal, vulkan, and webgpu overrides require their matching Cargo features and pin the graphics API. Explicit Metal selection requires native MSL support. The wgpu override keeps automatic API and compiler selection with WGSL fallback.

Panics when no execution backend is enabled, or when BURN_DEVICE names an unknown backend or one whose Cargo feature is not enabled.

impl Device for DispatchDevice

fn from_id(device_id: DeviceId) -> DispatchDevice

Create a device from its id.

fn to_id(&self) -> DeviceId

Retrieve the device id from the device.

impl DeviceOps for DispatchDevice

fn defaults(&self) -> DeviceSettings

Returns the default settings for this device.

fn id(&self) -> DeviceId

Returns the device id.

impl Eq for DispatchDevice

impl From<FlexDevice> for DispatchDevice

Available on crate feature flex only.

fn from(device: FlexDevice) -> DispatchDevice

Converts to this type from the input type.

impl PartialEq for DispatchDevice

fn eq(&self, other: &DispatchDevice) -> 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.

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

Tests for !=. The default implementation is almost always sufficient, and should not be overridden without very good reason.