[🤖] Dispatch: Int/Bool results of autodiff-tracked ops report a plain device, then silently re-lift to autodiff on `.float()`

**Describe the bug** On the Dispatch backend, taking an autodiff-tracked float and running a kind-changing op followed by a cast back to float, e.g. `x.argmax(1).float()` or `x.greater_equal_elem(1.0).float()` results in the tensor being promoted back to autodiff, where: - The Int/Bool intermediate reports `device().is_autodiff() == false`, - Yet the `.float()` result is a genuinely autodiff-tracked tensor again. The tensor's reported device and its actual routing disagree, so code that checks `is_autodiff()` to keep gradient-free work (metrics, logging) off the autodiff backend is silently defeated. Downstream this surfaces as *"tensors are not on the same backend"* panics when the recast result meets a genuinely plain tensor, or as unbounded autograd-node accumulation when it doesn't. **To Reproduce** <details> <summary>Reproduction</summary> - `int_into_float_relifts_to_autodiff` and `bool_into_float_relifts_to_autodiff` assert the *expected* behavior and **fail** on the pinned rev (the final assert sees the re-lifted autodiff tensor). - `plain_origin_does_not_lift` (control) passes: the same chain from a plain-device float stays plain. - `workaround_inner_before_the_cast_chain` passes: calling `.inner()` on the float *before* the `float→int/bool→float` chain strips the marker at the source, so the whole chain stays plain. ```toml [dependencies.burn] git = "https://github.com/tracel-ai/burn.git" rev = "d028234e2ccd75c5ca57d1c71ae74456e0e15c4f" # 2026-06-25 main default-features = false features = ["flex", "autodiff", "optim", "std"] ``` ```rs #[cfg(test)] mod autodiff_implicit_recast { use burn::prelude::*; fn ad_device() -> Device { Device::default().autodiff() } fn sample(device: &Device) -> Tensor<2> { Tensor::<2>::from_floats([[0.0f32, 1.0], [2.0, 0.5]], device) } #[test] fn int_into_float_relifts_to_autodiff() { let x = sample(&ad_device()); // (autodiff) assert!(x.device().is_autodiff()); // ok (autodiff) let x_idx = x.argmax(1); // float->int assert!(!x_idx.device().is_autodiff()); // ok (plain) let x_idx = x_idx.float(); // int->float assert!(!x_idx.device().is_autodiff(), "lifted to autodiff"); // panic (autodiff) } #[test] fn bool_into_float_relifts_to_autodiff() { let x = sample(&ad_device()); // (autodiff) assert!(x.device().is_autodiff()); // ok (autodiff) let x_mask = x.greater_equal_elem(1.0); // float->bool assert!(!x_mask.device().is_autodiff()); // ok (plain) let x_mask = x_mask.float(); // bool->float assert!(!x_mask.device().is_autodiff(), "lifted to autodiff"); // panic (autodiff) } /// This test passes. #[test] fn plain_origin_does_not_lift() { let x = sample(&Device::default()); // (plain) assert!(!x.device().is_autodiff()); // ok (plain) let x_idx = x.argmax(1); // float->int (plain) assert!(!x_idx.device().is_autodiff()); // ok (plain) let x_idx = x_idx.float(); // int->float assert!(!x_idx.device().is_autodiff()); // ok (plain) } /// USER-LEVEL FIX (this test passes): strip the float with `.inner()` BEFORE the /// float→int/bool→float chain — the plain-origin intermediates then carry no hidden /// autodiff marker, so the cast back to float stays plain. #[test] fn workaround_inner_before_the_cast_chain() { let x = sample(&ad_device()); // (autodiff) let x = x.inner(); // strip to the inner backend first assert!(!x.device().is_autodiff()); // ok (plain) let x_idx = x.clone().argmax(1).float(); // float->int->float, all plain assert!(!x_idx.device().is_autodiff()); // ok (plain) let x_mask = x.greater_equal_elem(1.0).float(); // float->bool->float, all plain assert!(!x_mask.device().is_autodiff()); // ok (plain) } } ``` </details> **Expected behavior** The reported device and the routing should agree: either Int/Bool results of autodiff-tracked ops report an autodiff device, or they drop the hidden autodiff marker so a later `.float()` stays on the plain backend — but a tensor whose `is_autodiff()` is `false` should never produce autodiff-tracked results. **Additional context** - _(Analysis and reproduction prepared by Claude)._ - An autodiff-dispatched op returning Int/Bool goes through `DispatchKindConversion<Autodiff<B, C>>::from_backend` (`burn-dispatch/src/tensor.rs`): the result's `kind` is a **plain** backend variant ("Pass-throughs for non-differentiable types"), but the returned `DispatchTensor` still gets **`checkpointing: Some(C::STRATEGY)`** — the autodiff marker survives on a plain-kinded tensor. - `DispatchTensor::device()` derives the device from `kind` only, so the Int/Bool tensor reports a plain device; the marker is folded into the device only when the kind is already autodiff. (There is a `TODO: should int and bool kinds return an autodiff device?` right there.) - A float-returning op on that tensor (`int_into_float` / `bool_into_float`, `unary_op_arms!(Float, ...)` in `burn-dispatch/src/macros.rs`) checks `if checkpointing.is_some()` and routes through `with_autodiff_backend!`, wrapping the result as a full autodiff float — the hidden marker becomes a tracked tensor again. - Note `.detach()` / `.set_require_grad(false)` on the intermediate don't help: they act on the graph/grad requirement, not on the carried `checkpointing` marker. - Linking to a discord [thread](https://discord.com/channels/1038839012602941528/1519569109904396459).