tracel-ai/burn · error

backend extension input is on the wrong backend

Error message

backend extension input is on the wrong backend

What it means

Emitted by generated code from extract_extension when mapping a dispatch tensor into a backend-extension primitive via `map_from_dispatch`. Inside the mapping, the inner dispatch kind must be the specific backend's variant; the `_ =>` arm panics when the extension input tensor belongs to a different backend than the one the extension is instantiated for. Extensions are per-backend code, so there is no fallback conversion.

Source

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

        }
    }
}

fn extract_extension(
    name: &syn::Ident,
    ty: &syn::Type,
    paths: &RoutingPaths,
    backend: &syn::Ident,
    autodiff: bool,
) -> TokenStream {
    let backend_root = &paths.backend_root;
    let dispatch_root = &paths.dispatch_root;
    let dispatch_kind = quote!(#dispatch_root::DispatchTensorKind);
    let backend_tensor = quote!(#dispatch_root::BackendTensor);
    let autodiff_trait = &paths.autodiff_trait;
    let extension_trait = &paths.extension_trait;
    let backend_alias = &paths.backend_alias;
    let mismatch = quote!(panic!("backend extension input is on the wrong backend"));
    if autodiff {
        let target = ir::with_backend(ty, quote!(#backend_alias));
        let context = quote!(#dispatch_root::DispatchAutodiffContext);
        quote! {
            let #name = <#target as #extension_trait<#backend_alias>>::map_from_dispatch(#name, |__tensor| {
                let __input_context = __tensor.autodiff;
                match __tensor.kind {
                    #dispatch_kind::Autodiff(inner) => {
                        let #context::Enabled(_) = __input_context else {
                            panic!("an autodiff float primitive must have an enabled autodiff context")
                        };
                        let tensor = match *inner {
                            #dispatch_kind::#backend(tensor) => tensor,
                            #[allow(unreachable_patterns)]
                            _ => #mismatch,
                        };
                        match tensor {
                            #backend_tensor::Autodiff(tensor) => #backend_tensor::Float(tensor),

View on GitHub (pinned to d16f7ba2ed)

Solutions

1. Instantiate the extension with the same backend that produced the input tensor
2. Convert the tensor onto the extension's backend before calling the extension op
3. Check that the `#backend_alias` used in your route matches your tensor construction site

Example fix

// before
let t = <Wgpu>::from_data(data, &device);
my_ext_op::<CandleBackedExt>(t); // panics: extension input on wrong backend
// after
let t = <Candle>::from_data(data, &device);
my_ext_op::<CandleBackedExt>(t);

if !matches!(&t.kind, DispatchTensorKind::ExtBackend(_) | DispatchTensorKind::Autodiff(_)) { panic!("extension requires tensors on its backend"); }

fn matches_extension_backend(t: &DispatchTensor) -> bool { matches!(&t.kind, DispatchTensorKind::ExtBackend(_)) }

When it happens

Trigger: Calling an extension op (e.g. a custom fused kernel or extended primitive routed through `DispatchAutodiffContext`) with a tensor whose inner `DispatchTensorKind` is not the extension backend's variant, in either the autodiff or non-autodiff path.

Common situations: Using a backend-specific extension (e.g. a CubeCL-only kernel) with tensors from the generic/default backend; wrapping a module for one backend but feeding it tensors from another after a refactor.

Related errors

• Tensor is on the wrong backend (expected {backend}).
• Autodiff float tensor is on the wrong backend (expected {bac
• q_matmul inputs are on different backends
• unsupported dispatch backend
• only float tensors may use an autodiff primitive

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