tracel-ai/burn · error

Autodiff not supported for custom op `{}`

Error message

Autodiff not supported for custom op `{}`

What it means

In burn-backend-extension's generated routing (crates/burn-backend-extension/src/routing.rs:1126), when a custom op was declared without autodiff support (autodiff=false / no autodiff route generated), the autodiff-context route asserts the context is Disabled; if autodiff is actually enabled (__ad_ctx says Enabled), it panics with unimplemented!("Autodiff not supported for custom op `{}`").

Source

Thrown at crates/burn-backend-extension/src/routing.rs:1126

        concrete_enabled,
        autodiff_call,
        concrete_enabled_output,
    } = fragments;
    let dispatch_root = &paths.dispatch_root;
    let context = quote!(#dispatch_root::DispatchAutodiffContext);
    let backend_root = &paths.backend_root;
    let backend_alias = &paths.backend_alias;
    let strategy = quote!(#dispatch_root::GradientCheckpointingStrategy);
    let attr = autodiff_attr.cloned().unwrap_or_default();
    let concrete_enabled_output = if operation.output.contains_float() {
        concrete_enabled_output
    } else {
        concrete_enabled
    };
    let concrete_route = if !autodiff {
        quote! {
            let #context::Disabled = __ad_ctx else {
                unimplemented!("Autodiff not supported for custom op `{}`", stringify!(#operation_name))
            };
            type #backend_alias = #backend_root::#backend;
            #concrete_inputs
            #concrete_invoke
            #concrete_disabled
        }
    } else if autodiff_attr.is_none() {
        quote! {
            type #backend_alias = #backend_root::#backend;
            #concrete_inputs
            #concrete_invoke
            match __ad_ctx {
                #context::Disabled => #concrete_disabled,
                #context::Enabled(__strategy) => #concrete_enabled_output,
            }
        }
    } else {
        quote! {

View on GitHub (pinned to d16f7ba2ed)

Solutions

1. Declare autodiff support for the custom op (add the autodiff attribute/route so an autodiff strategy is generated).
2. Run the op under no_grad / with autodiff disabled (Disabled context).
3. Use a builtin op with autodiff support for the training path.

Example fix

// before
#[extension(op)]
#[backend(...)] // no autodiff route
fn my_op(x: Tensor) -> Tensor { ... }
// used in training -> panic
// after
#[extension(op)]
#[backend(...)]
#[autodiff]
fn my_op(x: Tensor) -> Tensor { ... }

// only call the custom op outside autodiff if it was not declared with #[autodiff]
if grad_enabled && !op_declares_autodiff() { panic!("op requires autodiff support"); }

When it happens

Trigger: Calling a custom op declared without autodiff support while running under burn-autodiff with gradients enabled (Enabled context), e.g. inside a training loop.

Common situations: Using a custom op in a model's forward pass during training when the op only declared concrete (non-autodiff) backend routes.

Related errors

• Operation not marked for autodiff.
• an autodiff float primitive must have an enabled autodiff co
• Cannot move between autodiff and non-autodiff instances.
• Autodiff float tensor is on the wrong backend (expected {bac
• Expected autodiff-wrapped float tensor for backend {backend}

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