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

When dispatching a distributed float operation on an Autodiff tensor kind, the macro requires an enabled autodiff context (`DispatchAutodiffContext::Enabled`) carrying the checkpointing strategy. If the context is missing/disabled while the tensor is an autodiff float primitive, it panics.

Source

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

            dispatch_distributed_devices_arms,
            $device,
            $devices,
            |$inner_devices| $body
        )
    };
}

macro_rules! dispatch_distributed_float_arms {
    ($tensor:expr, |$inner:ident| $body:expr; $([$Backend:ident, $cfg:meta]),*) => {{
        let autodiff = $tensor.autodiff;
        match $tensor.kind {
            #[cfg(feature = "autodiff")]
            $crate::DispatchTensorKind::Autodiff(inner) => match *inner {
                $(
                    #[cfg($cfg)]
                    $crate::DispatchTensorKind::$Backend($inner) => {
                        let $crate::DispatchAutodiffContext::Enabled(checkpointing) = autodiff else {
                            panic!("an autodiff float primitive must have an enabled autodiff context")
                        };
                        with_autodiff_backend!($Backend, checkpointing, |B| {
                            let $inner = $inner.autodiff();
                            $crate::DispatchTensor {
                                kind: $crate::DispatchTensorKind::Autodiff(alloc::boxed::Box::new(
                                    $crate::DispatchTensorKind::$Backend(
                                        $crate::BackendTensor::Autodiff($body),
                                    ),
                                )),
                                autodiff,
                            }
                        })
                    }
                )*
                #[allow(unreachable_patterns)]
                other => panic!("Distributed operations are not supported for tensor kind {other:?}"),
            },
            $(

View on GitHub (pinned to d16f7ba2ed)

Solutions

  1. Ensure the autodiff context is Enabled (with a checkpointing/backward strategy) when dispatching ops on autodiff tensors
  2. Call the operation through the normal Autodiff backend API instead of the raw dispatch function
  3. Convert to the primitive tensor (`.primitive()`/detach) if gradients are not needed for that op

Example fix

// before
dispatch_fn(tensor, DispatchAutodiffContext::Disabled);
// after
dispatch_fn(tensor, DispatchAutodiffContext::Enabled(strategy));
Defensive patterns

Strategy: validation

Validate before calling

fn ensure_ad_enabled(ctx: &DispatchAutodiffContext) -> Result<(), &'static str> {
    match ctx { DispatchAutodiffContext::Enabled(_) => Ok(()), DispatchAutodiffContext::Disabled => Err("need enabled autodiff context") }
}

Prevention

When it happens

Trigger: Calling a distributed operation on an `AutodiffTensor` while the dispatch autodiff context is `Disabled` — e.g. running a distributed op through the autodiff backend without the autodiff feature providing a context, or calling the raw dispatch API without passing an Enabled context.

Common situations: Manual dispatch-layer calls (custom ops/serde graph execution) that forget to thread the autodiff context; running autodiff tensors under a code path compiled without proper context setup.

Related errors


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