tracel-ai/burn · error

an enabled float tensor must use an autodiff primitive

Error message

an enabled float tensor must use an autodiff primitive

What it means

When a tensor's autodiff context is Enabled but the tensor itself is a plain float primitive (not an `Autodiff` kind), `device()` panics: gradients cannot be tracked on a non-autodiff float primitive. The invariant is that an enabled float tensor must wrap an autodiff primitive; non-float tensors are permitted since they are just re-wrapped.

Source

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

    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) => {}
        }

        device
    }
}

impl DispatchTensorKind {
    /// Returns the backend tensor kind name.
    pub(crate) fn name(&self) -> &'static str {
        match self {

View on GitHub (pinned to d16f7ba2ed)

Solutions

1. Wrap the float primitive in the autodiff primitive (`DispatchTensorKind::Autodiff`) when the context is enabled.
2. Create the tensor via the autodiff-aware constructor so kind and context stay consistent.
3. Disable the autodiff context for tensors that are intentionally non-autodiff.

Example fix

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

if matches!(tensor.autodiff, DispatchAutodiffContext::Enabled(_)) {
    assert!(matches!(tensor.kind, DispatchTensorKind::Autodiff(_)), "enabled float tensor must be autodiff primitive");
}

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

let result = std::panic::catch_unwind(AssertUnwindSafe(|| tensor.device()));
match result {
    Ok(d) => use_device(d),
    Err(_) => eprintln!("enabled context on non-autodiff float primitive"),
}

When it happens

Trigger: Calling `device()` (e.g. via `assert_enabled_float`) on a tensor with `DispatchAutodiffContext::Enabled(_)` whose kind is a plain Float dispatch kind rather than `DispatchTensorKind::Autodiff`.

Common situations: Mixing tensors created inside and outside an autodiff session, custom backends wrapping raw float primitives while marking the context as autodiff-enabled, or partial migration from inference to training code paths.

Related errors

• an autodiff float primitive must have an enabled autodiff co
• 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/0b03090f5d7b908e. Report an issue: GitHub.