[🤖] `#[backend_extension]`: autodiff ops only work on the first backend in the list

_(Analysis and reproduction prepared by Claude)._ **Describe the bug** A backend extension whose op takes only float tensors can be called under autodiff on **just one** backend — the first one listed in the attribute. Any other listed backend panics with: ``` internal error: entered unreachable code: Autodiff backend mismatch ``` Cause: `gen_autodiff_arm` (`crates/burn-backend-extension/src/lib.rs`) emits one match arm per concrete backend, but the arm's pattern is keyed on the *tensor kind*, and every float argument yields the same pattern `DispatchTensorKind::Autodiff(..)`. All the arms are therefore identical, only the first is reachable, and its body unwraps that one backend: ```rust let #name = match *#name { burn::backend::DispatchTensorKind::#b_ident(t) => t.autodiff(), _ => unreachable!("Autodiff backend mismatch"), }; ``` Non-float arguments do get per-backend patterns, so an op with an int/bool tensor argument dispatches correctly — which is probably why this has gone unnoticed. Plain (non-autodiff) dispatch is unaffected. **To Reproduce** <details> <summary>Complete crate — CPU only, no GPU needed</summary> `Cargo.toml`: ```toml [package] name = "adrepro" version = "0.1.0" edition = "2024" [dependencies] burn = { git = "https://github.com/tracel-ai/burn.git", rev = "c2ecf30647ba9187b8e6ad54316895ed986c5234", default-features = false, features = [ "extension", "std", "autodiff", "flex", "ndarray", ] } ``` `src/main.rs`: ```rust use burn::backend::tensor::FloatTensor; // The macro names these unqualified, so they must be in scope. use burn::backend::{Autodiff, Backend, Dispatch, Flex, NdArray, backend_extension}; use burn::prelude::*; #[backend_extension(NdArray, Flex, Autodiff)] pub trait DoubleExt: Backend { fn double(x: FloatTensor<Self>) -> FloatTensor<Self> { Self::float_add(x.clone(), x) } } impl DoubleExt for NdArray {} impl DoubleExt for Flex {} impl<B: Backend + DoubleExt, C: burn::backend::autodiff::checkpoint::strategy::CheckpointStrategy> DoubleExt for Autodiff<B, C> { } fn run(label: &str, device: Device) { let x = Tensor::<1>::from_floats([1.0, 2.0, 3.0], &device); let y: Tensor<1> = Tensor::from_dispatch(<Dispatch as DoubleExt>::double(x.into_dispatch())); println!("{label}: {:?}", y.to_data().to_vec::<f32>().unwrap()); } fn main() { run("ndarray plain ", Device::ndarray()); run("flex plain ", Device::flex()); run("ndarray autodiff", Device::ndarray().autodiff()); // first listed: ok run("flex autodiff", Device::flex().autodiff()); // second listed: panics } ``` ```text ndarray plain : [2.0, 4.0, 6.0] flex plain : [2.0, 4.0, 6.0] ndarray autodiff: [2.0, 4.0, 6.0] thread 'main' panicked at src/main.rs:6:1: internal error: entered unreachable code: Autodiff backend mismatch ``` Swapping the attribute to `#[backend_extension(Flex, NdArray, Autodiff)]` moves the panic to `ndarray autodiff`, confirming it is purely list order. </details> **Expected behavior** The autodiff arm should dispatch on the backend actually inside `DispatchTensorKind::Autodiff(..)`, so every listed backend works. A binary compiled with several backends should be able to train on any of them, the same way it can already run non-autodiff ops on any of them. **Additional context** - The practical cost is that **one build cannot serve several backends**. Since the device is a runtime choice (`BURN_DEVICE`), a single binary would otherwise be enough to run the same workload on CPU and GPU — useful for benchmarking, CI matrices, and shipping one artifact. Today any workload that trains through a backend extension needs a separate build per backend, with the extension's attribute order matching the backend being run. - Version: `burn` git rev `c2ecf30647ba9187b8e6ad54316895ed986c5234`.