tracel-ai/burn · error

input tensor `{}` is on the wrong backend

Error message

input tensor `{}` is on the wrong backend

What it means

Generated by extract_tensor (borrowed float + autodiff path): the input tensor's kind is Autodiff, but its inner device does not match the concrete backend the operation was expanded for, so the macro panics naming the tensor. The dispatch macro generates per-backend code; at runtime the tensor's actual inner backend must equal that backend.

Source

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

            #dispatch_kind::Autodiff(_) => panic!("autodiff float input reached concrete dispatch"),
        }
    });
    let invalid_autodiff_kind = extraction.has_autodiff_variant.then(|| quote! {
        #ad_cfg
        #dispatch_kind::Autodiff(_) => panic!("only float tensors may use an autodiff primitive"),
    });

    if kind == TensorKind::Float && extraction.autodiff {
        if borrowed {
            let lifted = format_ident!("__lifted_{name}");
            quote! {
                let #lifted;
                let #name = match &#name.kind {
                    #dispatch_kind::Autodiff(__inner) => {
                        #validate_enabled_float
                        match __inner.as_ref() {
                            #dispatch_kind::#backend(__inner) => __inner.as_autodiff(),
                            _ => panic!("input tensor `{}` is on the wrong backend", stringify!(#name)),
                        }
                    }
                    #dispatch_kind::#backend(__inner) => {
                        #validate_disabled_float
                        #lifted = <#backend_alias as #autodiff_trait>::from_inner(__inner.as_float().clone());
                        &#lifted
                    }
                    _ => panic!("input tensor `{}` is on the wrong backend", stringify!(#name)),
                };
            }
        } else {
            quote! {
                let #name = match #name.kind {
                    #dispatch_kind::Autodiff(__inner) => {
                        #validate_enabled_float
                        match *__inner {
                            #dispatch_kind::#backend(__inner) => __inner.autodiff(),
                            _ => panic!("input tensor `{}` is on the wrong backend", stringify!(#name)),

View on GitHub (pinned to d16f7ba2ed)

Solutions

1. Ensure all tensor arguments of an operation come from the same backend/device as the dispatch target
2. Move/convert the tensor to the expected backend before the call (re-create it on the target device)
3. Check which backend the failing op was expanded for and align tensor creation accordingly

Example fix

// before
let a = Tensor::<Autodiff<Cpu>, 2>::ones(...);
let b = Tensor::<Autodiff<Wgpu>, 2>::ones(...);
mix(a, b); // panic
// after
let b = Tensor::<Autodiff<Cpu>, 2>::ones(...); // same backend as a
mix(a, b);

fn ensure_same_backend(a: &DispatchTensor, b: &DispatchTensor) -> Result<(), String> {
    if std::mem::discriminant(&a.kind) != std::mem::discriminant(&b.kind) {
        return Err("input tensors are on different backends".into());
    }
    Ok(())
}

fn inner_is_expected_backend(t: &DispatchTensor) -> bool {
    match &t.kind {
        DispatchTensorKind::Autodiff(inner) => matches!(inner.as_ref(), DispatchTensorKind::Cpu(_)),
        DispatchTensorKind::Cpu(_) => true,
        _ => false,
    }
}

let out = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| mix(&a, &b)))
    .map_err(|_| "wrong backend for input tensor; recreate it on the dispatch device".to_string())?

When it happens

Trigger: Calling a dispatch operation with a tensor whose inner backend differs from the backend the call was dispatched/compiled to — e.g. a Cpu tensor passed to a Wgpu-targeted expansion, wrapped in Autodiff.

Common situations: Mixing tensors created on different backend devices in one op; moving a model/tensor between backends without re-dispatching; multi-backend setups (cpu + wgpu + cuda) where tensors come from the wrong device.

Related errors

• Can't differentiate avg pool 2d backward.
• Can't differentiate max pool2d with indices backward.
• Can't differentiate adaptive avg pool2d backward.
• Can't differentiate adaptive avg pool3d backward.
• Can't differentiate interpolate backward.

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