tracel-ai/burn · error

an autodiff float primitive must have an enabled autodiff co

Error message

an autodiff float primitive must have an enabled autodiff context

What it means

During `device()` resolution, a tensor whose kind is `DispatchTensorKind::Autodiff` must carry `DispatchAutodiffContext::Enabled`. If the autodiff context is `Disabled` while the kind is autodiff, the internal invariant is broken and the code panics. This keeps the kind and the autodiff context consistent so device extraction can find the autodiff device.

Source

Thrown at crates/burn-dispatch/src/tensor.rs:388

    fn shape(&self) -> Shape {
        self.kind.shape()
    }

    fn can_mut(&self) -> bool {
        self.kind.can_mut()
    }

    type Device = DispatchDevice;

    fn device(&self) -> Self::Device {
        #[allow(unused_mut)]
        let mut device = self.kind.device();

        #[cfg(feature = "autodiff")]
        match (&self.kind, self.autodiff) {
            (DispatchTensorKind::Autodiff(_), DispatchAutodiffContext::Disabled) => {
                panic!("an autodiff float primitive must have an enabled autodiff context")
            }
            (DispatchTensorKind::Autodiff(_), DispatchAutodiffContext::Enabled(strategy)) => {
                let DispatchDevice::Autodiff(device) = &mut device else {
                    unreachable!("autodiff primitive must report an autodiff device")
                };
                device.checkpointing = strategy;
            }
            (_, DispatchAutodiffContext::Enabled(strategy)) => {
                if self.dtype().is_float() {
                    panic!("an enabled float tensor must use an autodiff primitive")
                }
                device = DispatchDevice::autodiff(device);
                let DispatchDevice::Autodiff(device) = &mut device else {
                    unreachable!()
                };
                device.checkpointing = strategy;
            }
            (_, DispatchAutodiffContext::Disabled) => {}

View on GitHub (pinned to d16f7ba2ed)

Solutions

1. Ensure the autodiff context is Enabled (with a strategy) whenever the tensor kind is Autodiff.
2. Fix the constructor/transform that produced the inconsistent kind/context pair.
3. If autodiff is intentionally disabled, rebuild the tensor with a non-autodiff kind.

Example fix

// before
DispatchTensor { kind: DispatchTensorKind::Autodiff(p), autodiff: DispatchAutodiffContext::Disabled }
// after
DispatchTensor { kind: DispatchTensorKind::Autodiff(p), autodiff: DispatchAutodiffContext::Enabled(strategy) }

assert!(matches!(tensor.autodiff, DispatchAutodiffContext::Enabled(_)), "autodiff kind requires enabled context");

fn has_enabled_autodiff(t: &DispatchTensor) -> bool {
    !matches!(t.kind, DispatchTensorKind::Autodiff(_))
        || matches!(t.autodiff, DispatchAutodiffContext::Enabled(_))
}

let result = std::panic::catch_unwind(AssertUnwindSafe(|| tensor.device()));
match result {
    Ok(d) => use_device(d),
    Err(_) => eprintln!("autodiff tensor had disabled context"),
}

When it happens

Trigger: Constructing or transforming a `DispatchTensor` so that `kind` is `Autodiff(_)` while `autodiff` is `DispatchAutodiffContext::Disabled`, then calling `device()` (e.g. via `assert_enabled_float`).

Common situations: Manual construction of DispatchTensors, custom backends/wrappers that set the kind and autodiff context independently, or code paths that disable autodiff without retagging the tensor kind.

Related errors

• an enabled float tensor must use an autodiff primitive
• Autodiff should not wrap an autodiff tensor.
• Requires autodiff tensor.
• an enabled float tensor must use an autodiff primitive
• Should be float, got autodiff

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