What causes "Autodiff should not wrap an autodiff device." in tracel-ai/burn?

Unwrap the inner device and pass the concrete backend device instead of double-wrapping

Autodiff should not wrap an autodiff device.

Error message

Autodiff should not wrap an autodiff device.

What it means

The distributed device dispatch macro explicitly rejects `DispatchDevice::Autodiff`: autodiff is a wrapper around a backend, not a physical device, so nesting it is always a programming error and the library panics with 'Autodiff should not wrap an autodiff device'.

Source

Thrown at crates/burn-dispatch/src/ops/distributed.rs:97

                $crate::DispatchDevice::$Backend(_) => {
                    type B = $crate::backends::Autodiff<$crate::backends::$Backend>;
                    let $inner_devices = $devices
                        .iter()
                        .map(|d| {
                            // In a build with only one dispatch variant compiled in, the
                            // pattern is irrefutable and the else clause is dead — which is
                            // the point: it only guards the multi-backend builds.
                            #[allow(irrefutable_let_patterns)]
                            let DispatchDevice::$Backend(dev) = d else {
                                unreachable!("All devices are expected to be of the same variant.")
                            };
                            dev.clone()
                        })
                        .collect::<Vec<_>>();
                    $body
                }
            )*
            $crate::DispatchDevice::Autodiff(_) => panic!("Autodiff should not wrap an autodiff device."),
            #[allow(unreachable_patterns)]
            other => panic!("Distributed operations are not supported for device {other:?}"),
        }
    };
}

/// Dispatches an operation body based on the provided devices.
macro_rules! dispatch_distributed_devices {
    ($device:expr, $devices:expr, |$inner_devices:ident| $body:expr) => {
        distributed_backend_list!(
            dispatch_distributed_devices_arms,
            $device,
            $devices,
            |$inner_devices| $body
        )
    };
}

View on GitHub (pinned to d16f7ba2ed)

Solutions

Unwrap the inner device and pass the concrete backend device instead of double-wrapping
When initializing distributed workers, use the base backend's device, not the Autodiff<B> device
Guard wrapping logic so it only applies Autodiff once

Example fix

// before
let dev = DispatchDevice::Autodiff(Box::new(ad_device)); // double wrap
// after
let dev = ad_device.base_device(); // pass the underlying backend device

Defensive patterns

Strategy: validation

Validate before calling

fn unwrap_ad_device(dev: &DispatchDevice) -> DispatchDevice {
    match dev { DispatchDevice::Autodiff(inner) => unwrap_ad_device(inner), d => d.clone() }
}

Type guard

fn is_base_device(d: &DispatchDevice) -> bool { !matches!(d, DispatchDevice::Autodiff(_)) }

Prevention

Never wrap Autodiff twice; use base_backend()/inner device accessors
In device-building helpers, unwrap to the concrete backend device
Assert with `matches!` that a device is a concrete backend before distributed init

When it happens

Trigger: Constructing or passing a DispatchDevice::Autodiff that itself contains another Autodiff device — e.g. wrapping an already-autodiff-wrapped device when setting up a distributed operation or distributed parameters.

Common situations: Manually building dispatch devices for distributed training and accidentally double-wrapping with Autodiff; helper code that wraps devices unconditionally.

Related errors

AI-assisted analysis of tracel-ai/burn@d16f7ba2ed (2026-09-05). Data as JSON: /api/errors/88aeb885ca190516. Report an issue: GitHub.