tracel-ai/burn · error

Operation not marked for autodiff.

Error message

Operation not marked for autodiff.

What it means

A tensor whose DispatchTensorKind is Autodiff was passed to a device-movement/matching path without a corresponding autodiff arm enabled — the macro's autodiff-kind catch-all panics 'Operation not marked for autodiff.' This means the operation or macro invocation was generated without the autodiff marker even though the tensor carries autodiff state, so the dispatcher cannot know how to propagate gradients or move the tensor correctly.

Source

Thrown at crates/burn-dispatch/src/macros.rs:230

                        type B2 = $crate::backends::$B2;
                        let $inner = tensor.$inner_fn();

                        $crate::DispatchTensor {
                            kind: $crate::DispatchTensorKind::$B2(
                                $crate::BackendTensor::$kind($body)
                            ),
                            autodiff: $crate::DispatchAutodiffContext::Enabled(
                                device_ad.checkpointing,
                            ),
                        }

                    },
                )+
            )*
            #[cfg(feature = "autodiff")]
            (_, $crate::DispatchDevice::Autodiff(_)) => unreachable!("Autodiff should not wrap an autodiff device."),
            #[cfg(feature = "autodiff")]
            ($crate::DispatchTensorKind::Autodiff(..), _) => panic!("Operation not marked for autodiff."),
            // Capture is intentionally one-way: initialized values can be moved onto a
            // capture device, but captured tensors have no materialized data to move back.
            #[cfg(feature = "capture")]
            ($crate::DispatchTensorKind::Capture(_), _) => {
                panic!("Cannot move a tensor from a capture device")
            }
        }
    };
}

/// Handles tensor movement between devices, supporting both same-backend transfers
/// and cross-backend dispatches.
macro_rules! to_device {
    ($kind:ident, $inner_fn:ident, $tensor:expr, $device:expr, $to_device:ident, |$inner:ident, $device_ident:ident| $body:expr) => {
        backend_matrix!(
            to_device_arms,
            $kind,
            $inner_fn,

View on GitHub (pinned to d16f7ba2ed)

Solutions

  1. Register the operation for autodiff using burn_autodiff's operation macros (e.g. `use burn_autodiff::ops::*` and the corresponding backward-op registration) so an autodiff arm is generated.
  2. If the op should not differentiate, convert the tensor to a non-autodiff kind (detach via `.inner()` / backend tensor) before the call.
  3. Ensure the `autodiff` feature is enabled consistently across all crates in the workspace so generated arms match.
  4. Wrap the call in the dispatch macro variant intended for autodiff (marked arms), not the plain backend variant.

Example fix

// before
tensor.dispatch(op) // Autodiff-kind tensor, op has no autodiff arm -> panic

// after
use burn_autodiff::ops::register_op;
register_op!(MyOp); // generates the autodiff arm
tensor.dispatch(MyOp)
Defensive patterns

Strategy: type-guard

Validate before calling

fn ensure_marked_for_autodiff(kind: &DispatchTensorKind) {
    if matches!(kind, DispatchTensorKind::Autodiff(_)) {
        assert!(op_registered_for_autodiff(), "op must use burn_autodiff::ops registration");
    }
}

Type guard

fn is_autodiff_tensor(t: &DispatchPrimitive) -> bool {
    matches!(t.kind, DispatchTensorKind::Autodiff(_))
}

Prevention

When it happens

Trigger: Calling a backend operation on an Autodiff-kind tensor while the operation was not registered/generated with the autodiff arm (the `#[cfg(feature = "autodiff")]` kind arm at macros.rs:230 fires when the operation/tensor-kind pair has no marked match arm) — e.g. using a raw backend tensor op inside an Autodiff context, or a custom op not wrapped with burn_autodiff's registration macros.

Common situations: Implementing custom autodiff ops without `#[derive]`/macro registration; calling lower-level backend APIs directly with dispatch tensors; mixing non-autodiff operation macro instances with autodiff tensors after toggling the autodiff feature.

Related errors


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