fix!: address fusion and autodiff edge cases (#5400)

* Add common `GraphIr` and refactor burn-router to use it instead

* Add core ONNX export from existing `GraphIr`

* Add graph capture backend

* Add the capture device to burn-dispatch

* Add module ONNX export (static)

* Add dynamic module export and validated parameter bindings

* Cleanup dynamic axes & validation

* Add resnet18 fixture w/ test

* Lower batch norm to ONNX BatchNormalization

* Add interpolate

* Add FullOpIr -> ConstantOfShape support

* Add Full + SliceAssign -> ConstantPad and Concat op

* Add Neg op support

* Update dep to git

* Add OnnxModel w/ save

* boxblur to_device

* make it easier to pass a custom renderer

* add some funcs for grad checkpt

* replace inner() calls with no_grad()

* fix fusion panic

* fusion nhwc relayout fix + unit tests

* Remove stale onnx-export leftovers inherited from base branch

The branch was originally based on the onnx-export branch, whose crate
was since reworked and removed on main. Restore Cargo.toml/Cargo.lock
from main and drop crates/burn-onnx-export so the PR diff only contains
this branch's own changes.

* Remove leftover capture-device wiring from base branch

The CaptureDevice re-export and the Flex<->Capture backend_matrix rows
came from the pre-merge version of the capture backend on the old base
branch; the version merged on main (#5377) does not include them.

* de-claudify

* de-claudify

* renderer signature

* fix: address code review findings

Bound a cached plan's relative shape ids by what the plan's own operations
assigned rather than by the converter's live count. Plans usually fire from
`ExecutionTrigger::OnOperations`, i.e. exactly when later operations are
already queued behind them, and those inflated the count enough to let an
unfitting plan reach `OutputPlanner` and panic. `OperationQueue` now tracks
`shapes_assigned` alongside `relative`.

Also:
- Replace the leftover `[probe]` panic in `OutputPlanner` with a plain `expect`.
- Give `AutodiffDevice` a hand-written `PartialEq` comparing hardware identity
  only, so it agrees with `DispatchDevice`'s own `PartialEq` instead of
  contradicting it on the checkpointing strategy.
- Keep `default_renderer`'s parameters from warning as unused without `tui`.
- Move `BoxBlur`'s kernel on-device instead of round-tripping through the host.
- Drop the mnist example's local debugging `burn.toml`.

* Clarify gradient checkpointing API

---------

Co-authored-by: Guillaume Lagrange <lagrange.guillaume.1@gmail.com>
Co-authored-by: nathaniel <nathaniel.simard.42@gmail.com>

crates/burn-backend-extension/src/lib.rs

@@ -773,7 +773,7 @@ fn gen_tensor_input_dispatch_body(ir: &Extension, op: &Operation) -> TokenStream
         // Compute the checkpointing strategy and backend tag in a scoped block so the representative's
         // borrow of the inputs ends before the dispatch arms below move them.
         let (checkpointing, __burn_backend_tag): (
-            Option<burn::backend::CheckpointingStrategy>,
+            Option<burn::backend::GradientCheckpointingStrategy>,
             (bool, usize),
         ) = {
             let __repr: &burn::backend::DispatchTensor = (#float_chain)
@@ -1317,7 +1317,7 @@ pub fn derive_extension_type(input: TokenStream) -> TokenStream {
             fn map_to_dispatch<F>(
                 self,
                 map_kind: F,
-                checkpointing: Option<burn::backend::CheckpointingStrategy>,
+                checkpointing: Option<burn::backend::GradientCheckpointingStrategy>,
             ) -> Self::Target
             where
                 F: Fn(burn::backend::BackendTensor<B>) -> burn::backend::DispatchTensorKind,

crates/burn-backend-tests/tests/fusion/nhwc_relayout.rs

@@ -1,7 +1,7 @@
 use super::*;
 use burn_tensor::{
     Device, TensorData,
-    module::{interpolate, max_pool2d},
+    module::{adaptive_avg_pool2d, interpolate, max_pool2d},
     ops::{InterpolateMode, InterpolateOptions},
 };
 
@@ -173,6 +173,34 @@ fn fusion_test_nhwc_relayout_broadcast_different_shapes() {
     );
 }
 
+/// A tensor read by both the pool and a later operation must keep its NCHW layout.
+///
+/// The relayout rewrites where the pooled tensor's elements live so the pool can read it
+/// without a copy. That is only sound when the pool is that tensor's last reader.
+#[test]
+fn fusion_test_nhwc_relayout_shared_input_keeps_nchw_layout() {
+    let dev: Device = Default::default();
+
+    let input = TestTensor::<4>::from_data(
+        TensorData::from([[[[1.0f32, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]]]),
+        &dev,
+    );
+    dev.sync().unwrap();
+
+    // An elementwise op the fuser can absorb, whose result has two readers.
+    let filtered = input * 2.0;
+    let pooled = adaptive_avg_pool2d(filtered.clone(), [1, 1]);
+    let output = filtered * pooled;
+
+    // Doubled, the channel means are 5 and 13, so every product is exact in f32.
+    let expected = TensorData::from([[
+        [[10.0f32, 20.0], [30.0, 40.0]],
+        [[130.0, 156.0], [182.0, 208.0]],
+    ]]);
+
+    output.into_data().assert_eq(&expected, false);
+}
+
 fn seq_input(dev: &Device) -> TestTensor<4> {
     let data: Vec<f32> = (0..16).map(|i| i as f32).collect();
     TestTensor::from_data(TensorData::new(data, [2, 2, 2, 2]), dev)

crates/burn-core/src/backend.rs

@@ -39,7 +39,7 @@ pub trait ExtensionType<B: Backend> {
     fn map_to_dispatch<F>(
         self,
         map_kind: F,
-        checkpointing: Option<CheckpointingStrategy>,
+        checkpointing: Option<GradientCheckpointingStrategy>,
     ) -> Self::Target
     where
         F: Fn(BackendTensor<B>) -> DispatchTensorKind;

crates/burn-core/src/module/base.rs

@@ -585,6 +585,28 @@ mod tests {
         assert!(!module.weight.require_grad); // stateful
     }
 
+    /// `valid` on a module already on the inner backend returns it unchanged rather than
+    /// panicking.
+    #[test]
+    fn valid_is_idempotent() {
+        let device = test_device().autodiff();
+        let module = SimpleLinear::new(4, 4, &device).valid();
+
+        let module = module.valid();
+
+        assert!(!module.weight.is_require_grad());
+        assert!(!module.weight.val().device().is_autodiff());
+    }
+
+    #[test]
+    fn valid_on_a_plain_device_module_is_no_op() {
+        let module = SimpleLinear::new(4, 4, &test_device());
+
+        let module = module.valid();
+
+        assert!(!module.weight.val().device().is_autodiff());
+    }
+
     #[test]
     fn freeze_group_freezes_only_selected_params() {
         let device = test_device().autodiff();

crates/burn-core/src/module/param/constant.rs

@@ -135,7 +135,7 @@ impl<const D: usize, K: Basic> ModuleDisplay for Tensor<D, K> {}
 
 impl<const D: usize, K: Autodiff> AutodiffModule for Tensor<D, K> {
     fn valid(&self) -> Self {
-        self.clone().inner()
+        self.clone().no_grad()
     }
 
     fn from_inner(tensor: Self) -> Self {

crates/burn-core/src/module/param/running.rs

@@ -221,7 +221,7 @@ impl<const D: usize> AutodiffModule for RunningState<Tensor<D>> {
         self.sync();
         let value = self.value();
 
-        RunningState::with_id(self.id, value.inner())
+        RunningState::with_id(self.id, value.no_grad())
     }
 
     fn from_inner(module: Self) -> Self {

crates/burn-core/src/module/param/tensor.rs

@@ -317,7 +317,7 @@ impl<const D: usize> AutodiffModule for Param<Tensor<D>> {
         // Preserve initialized param `require_grad` state, but reset the inner value's.
         // `val()` folds any reparameterization into the base for inference.
         let require_grad = self.require_grad;
-        let mut param = Param::initialized(self.id, self.val().inner().set_require_grad(false));
+        let mut param = Param::initialized(self.id, self.val().no_grad().set_require_grad(false));
         param.require_grad = require_grad;
         param
     }
@@ -339,7 +339,7 @@ impl<const D: usize> AutodiffModule for Param<Tensor<D>> {
 
 impl<const D: usize> AutodiffModule for Param<Tensor<D, Int>> {
     fn valid(&self) -> Self {
-        Param::initialized(self.id, self.val().inner())
+        Param::initialized(self.id, self.val().no_grad())
     }
 
     fn from_inner(module: Self) -> Self {
@@ -349,7 +349,7 @@ impl<const D: usize> AutodiffModule for Param<Tensor<D, Int>> {
 
 impl<const D: usize> AutodiffModule for Param<Tensor<D, Bool>> {
     fn valid(&self) -> Self {
-        Param::initialized(self.id, self.val().inner())
+        Param::initialized(self.id, self.val().no_grad())
     }
 
     fn from_inner(module: Self) -> Self {

crates/burn-cubecl-fusion/src/engine/launch/output.rs

@@ -204,7 +204,9 @@ impl<'a, R: Runtime> OutputPlanner<'a, R> {
                             let pos = plan.runtime_layouts.len();
                             let mut shape_global = shape.clone();
                             for (i, s) in shape.iter().enumerate() {
-                                shape_global[i] = *context.shapes_relative2global.get(s).unwrap();
+                                shape_global[i] = *context.shapes_relative2global.get(s).expect(
+                                    "reference shape ids to be assigned by the running stream",
+                                );
                             }
 
                             let strides = strides_dyn_rank(&shape_global);

crates/burn-cubecl-fusion/src/engine/trace/base.rs

@@ -28,6 +28,16 @@ pub struct FuseTrace {
     pub resources: FuseResources,
 }
 
+impl FuseTrace {
+    /// The highest relative shape id any of this trace's blocks names as its reference shape.
+    pub fn max_relative_shape_id(&self) -> Option<usize> {
+        self.blocks
+            .iter()
+            .flat_map(|block| block.shape_ref.iter().copied())
+            .max()
+    }
+}
+
 impl core::fmt::Display for FuseTrace {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         writeln!(f, "FuseTrace")?;

crates/burn-cubecl-fusion/src/optim/base.rs

@@ -22,6 +22,12 @@ pub trait FusedOperation<R: Runtime>: Send + 'static {
     /// The number of operations fused.
     fn num_ops_fused(&self) -> usize;
 
+    /// The highest relative shape id this operation's traces name, if any. See
+    /// [`NumOperations::max_relative_shape_id`](burn_fusion::NumOperations::max_relative_shape_id).
+    fn max_relative_shape_id(&self) -> Option<usize> {
+        None
+    }
+
     /// Run the fused operation. `fallback` builds the unfused operation at
     /// the given index within the segment, for implementations that need to
     /// run part of it unfused (autotune fallbacks).
@@ -90,6 +96,10 @@ impl<R: Runtime> burn_fusion::NumOperations for CubeOptimization<R> {
         self.num_ops_fused()
     }
 
+    fn max_relative_shape_id(&self) -> Option<usize> {
+        self.optimization.max_relative_shape_id()
+    }
+
     fn name(&self) -> &'static str {
         Self::name(self)
     }
@@ -101,6 +111,7 @@ impl<R: Runtime> burn_fusion::NumOperations for CubeOptimization<R> {
 trait DynFusedOperation<R: Runtime>: Send {
     fn name(&self) -> &'static str;
     fn num_ops_fused(&self) -> usize;
+    fn max_relative_shape_id(&self) -> Option<usize>;
     fn run(
         &mut self,
         context: &mut Context<CubeFusionHandle<R>>,
@@ -118,6 +129,10 @@ impl<R: Runtime, O: FusedOperation<R>> DynFusedOperation<R> for O {
         FusedOperation::num_ops_fused(self)
     }
 
+    fn max_relative_shape_id(&self) -> Option<usize> {
+        FusedOperation::max_relative_shape_id(self)
+    }
+
     fn run(
         &mut self,
         context: &mut Context<CubeFusionHandle<R>>,

crates/burn-cubecl-fusion/src/optim/elemwise/optimization.rs

@@ -145,6 +145,10 @@ fn elemwise_fuse(
 pub const NAME: &str = "ElementWise";
 
 impl<R: Runtime> FusedOperation<R> for ElemwiseOptimization<R> {
+    fn max_relative_shape_id(&self) -> Option<usize> {
+        self.trace.max_relative_shape_id()
+    }
+
     const NAME: &'static str = self::NAME;
     type State = ElemwiseOptimizationState;
 

crates/burn-cubecl-fusion/src/optim/matmul/optimization.rs

@@ -712,6 +712,16 @@ fn launch_inner_fix_dtype<R: Runtime, A: BatchMatmulRoutine<()>>(
 pub const NAME: &str = "Matmul";
 
 impl<R: Runtime> FusedOperation<R> for MatmulOptimization<R> {
+    fn max_relative_shape_id(&self) -> Option<usize> {
+        [
+            self.info.trace.max_relative_shape_id(),
+            self.info.trace_fallback.max_relative_shape_id(),
+        ]
+        .into_iter()
+        .flatten()
+        .max()
+    }
+
     const NAME: &'static str = self::NAME;
     type State = MatmulOptimizationState;
 

crates/burn-cubecl-fusion/src/optim/nhwc_relayout/fuser.rs

@@ -6,7 +6,7 @@ use crate::{
     optim::{CubeOptimization, nhwc_relayout::optimization::NHWCRelayoutOptimization},
 };
 use burn_fusion::{FuserProperties, FuserStatus, OperationFuser};
-use burn_ir::{ModuleOperationIr, OperationIr, TensorIr};
+use burn_ir::{ModuleOperationIr, OperationIr, TensorIr, TensorStatus};
 use burn_std::Shape;
 use cubecl::Runtime;
 
@@ -88,7 +88,12 @@ impl<R: Runtime> OperationFuser<CubeOptimization<R>> for NHWCRelayoutFuser<R> {
         }
 
         match operation {
-            OperationIr::Module(ir) if let Some(tensor) = nhwc_relayout_tensor(ir) => {
+            // A `ReadOnly` tensor cannot be a relayout candidate: another consumer
+            // is still expecting NCHW.
+            OperationIr::Module(ir)
+                if let Some(tensor) = nhwc_relayout_tensor(ir)
+                    && tensor.status == TensorStatus::ReadWrite =>
+            {
                 self.op = Some(operation.clone());
                 self.fuser
                     .output_nhwc_layout(tensor, permutation(tensor.shape.num_dims()));

crates/burn-cubecl-fusion/src/optim/nhwc_relayout/optimization.rs

@@ -67,6 +67,10 @@ impl<R: Runtime> NHWCRelayoutOptimization<R> {
 pub const NAME: &str = "NHWCRelayout";
 
 impl<R: Runtime> FusedOperation<R> for NHWCRelayoutOptimization<R> {
+    fn max_relative_shape_id(&self) -> Option<usize> {
+        self.trace.max_relative_shape_id()
+    }
+
     const NAME: &'static str = self::NAME;
     type State = RelayoutOptimizationState;
 

crates/burn-cubecl-fusion/src/optim/reduce/optimization.rs

@@ -530,7 +530,25 @@ pub fn reduce_kernel_fused<In: Numeric, SizeIn: Size, Out: Numeric, SizeOut: Siz
 /// Name of the reduce fusion optimization.
 pub const NAME: &str = "Reduce";
 
+impl<R: Runtime> ReduceOptimizationInfo<R> {
+    /// The highest relative shape id across this reduce's trace and both of its fallbacks.
+    pub(crate) fn max_relative_shape_id(&self) -> Option<usize> {
+        [
+            self.trace.max_relative_shape_id(),
+            self.trace_read_fallback.max_relative_shape_id(),
+            self.trace_write_fallback.max_relative_shape_id(),
+        ]
+        .into_iter()
+        .flatten()
+        .max()
+    }
+}
+
 impl<R: Runtime> FusedOperation<R> for ReduceOptimization<R> {
+    fn max_relative_shape_id(&self) -> Option<usize> {
+        self.info.max_relative_shape_id()
+    }
+
     const NAME: &'static str = self::NAME;
     type State = ReduceOptimizationState;
 

crates/burn-cubecl-fusion/src/optim/reduce_broadcasted/optimization.rs

@@ -223,6 +223,27 @@ impl<R: Runtime> ReduceBroadcastedOptimization<R> {
 pub const NAME: &str = "ReduceBroadcasted";
 
 impl<R: Runtime> FusedOperation<R> for ReduceBroadcastedOptimization<R> {
+    fn max_relative_shape_id(&self) -> Option<usize> {
+        let fallbacks = self
+            .info
+            .fallbacks
+            .iter()
+            .filter_map(|fallback| match fallback {
+                ReduceBlockOptimInfo::Reduce(info) => info.max_relative_shape_id(),
+                ReduceBlockOptimInfo::Elemwise(opt) => {
+                    FusedOperation::<R>::max_relative_shape_id(opt.as_ref())
+                }
+            });
+
+        self.info
+            .broadcasted
+            .trace
+            .max_relative_shape_id()
+            .into_iter()
+            .chain(fallbacks)
+            .max()
+    }
+
     const NAME: &'static str = self::NAME;
     type State = ReduceBroadcastedOptimizationState;
 

crates/burn-dispatch/src/backend.rs

@@ -708,7 +708,7 @@ impl AutodiffBackend for Dispatch {
         let checkpointing = if let Some(strategy) = checkpointing {
             Some(strategy)
         } else {
-            Some(crate::CheckpointingStrategy::None)
+            Some(crate::GradientCheckpointingStrategy::Disabled)
         };
         DispatchTensor {
             kind,
@@ -773,7 +773,7 @@ impl AutodiffBackend for Dispatch {
         let checkpointing = if let Some(strategy) = checkpointing {
             Some(strategy)
         } else {
-            Some(crate::CheckpointingStrategy::None)
+            Some(crate::GradientCheckpointingStrategy::Disabled)
         };
         DispatchTensor {
             kind,

crates/burn-dispatch/src/device.rs

@@ -137,15 +137,36 @@ impl DispatchDevice {
 /// A wrapper that enables automatic differentiation for a [`DispatchDevice`].
 ///
 /// Use [`DispatchDevice::autodiff`] to construct this type.
-#[derive(Debug, Clone, PartialEq, Eq)]
+#[derive(Debug, Clone)]
 pub struct AutodiffDevice {
     pub(crate) inner: Box<DispatchDevice>,
-    pub(crate) checkpointing: CheckpointingStrategy,
+    pub(crate) checkpointing: GradientCheckpointingStrategy,
+}
+
+/// Compares on hardware identity only, ignoring the checkpointing strategy, so that this agrees
+/// with [`DispatchDevice`]'s own [`PartialEq`] — which has to ignore it, since comparing an
+/// `Autodiff` device against a raw one has no strategy to compare against. A derived impl would
+/// make `Autodiff(a) == Autodiff(b)` disagree with `DispatchDevice::Autodiff(a) ==
+/// DispatchDevice::Autodiff(b)`.
+///
+/// Use [`gradient_checkpointing_strategy`](Self::gradient_checkpointing_strategy) when the
+/// strategy is what you actually need to compare.
+#[cfg(feature = "autodiff")]
+impl PartialEq for AutodiffDevice {
+    fn eq(&self, other: &Self) -> bool {
+        self.inner == other.inner
+    }
 }
 
+#[cfg(feature = "autodiff")]
+impl Eq for AutodiffDevice {}
+
 #[cfg(feature = "autodiff")]
 impl AutodiffDevice {
-    pub(crate) fn new(device: DispatchDevice, checkpointing: CheckpointingStrategy) -> Self {
+    pub(crate) fn new(
+        device: DispatchDevice,
+        checkpointing: GradientCheckpointingStrategy,
+    ) -> Self {
         Self {
             inner: Box::new(device),
             checkpointing,
@@ -156,6 +177,11 @@ impl AutodiffDevice {
     pub fn inner(self) -> DispatchDevice {
         *self.inner
     }
+
+    /// Returns the gradient checkpointing strategy.
+    pub fn gradient_checkpointing_strategy(&self) -> GradientCheckpointingStrategy {
+        self.checkpointing
+    }
 }
 
 #[cfg(feature = "autodiff")]
@@ -170,24 +196,26 @@ impl core::ops::Deref for AutodiffDevice {
 
 #[allow(missing_docs)]
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
-/// Checkpointing strategy for autodiff.
+/// Gradient checkpointing strategy for autodiff.
 #[repr(u8)]
-pub enum CheckpointingStrategy {
+pub enum GradientCheckpointingStrategy {
+    /// Recompute selected activations during backpropagation to reduce peak memory usage.
     Balanced,
+    /// Disable gradient checkpointing while retaining autodiff tracking.
     #[default]
-    None,
+    Disabled,
 }
 
 #[cfg(feature = "autodiff")]
 pub(crate) fn validate_checkpointing(
-    lhs: Option<crate::CheckpointingStrategy>,
-    rhs: Option<crate::CheckpointingStrategy>,
-) -> Option<crate::CheckpointingStrategy> {
+    lhs: Option<crate::GradientCheckpointingStrategy>,
+    rhs: Option<crate::GradientCheckpointingStrategy>,
+) -> Option<crate::GradientCheckpointingStrategy> {
     match (lhs, rhs) {
         (Some(lhs), Some(rhs)) => {
             assert_eq!(
                 lhs, rhs,
-                "Autodiff strategy mismatch: {lhs:?} vs {rhs:?}. Tensors in the same operation must share a strategy."
+                "Gradient checkpointing strategy mismatch: {lhs:?} vs {rhs:?}. Tensors in the same operation must share a strategy."
             );
             Some(lhs)
         }
@@ -229,7 +257,11 @@ impl core::fmt::Debug for DispatchDevice {
             Self::Capture(device) => f.debug_tuple("Capture").field(device).finish(),
             #[cfg(feature = "autodiff")]
             // Format without `AutodiffDevice` wrapper
-            Self::Autodiff(device) => f.debug_tuple("Autodiff").field(&device.inner).finish(),
+            Self::Autodiff(device) => f
+                .debug_struct("Autodiff")
+                .field("device", &device.inner)
+                .field("checkpointing", &device.checkpointing)
+                .finish(),
         }
     }
 }
@@ -371,7 +403,9 @@ impl PartialEq for DispatchDevice {
         match (self, other) {
             // If both are Autodiff, compare the inner devices
             #[cfg(feature = "autodiff")]
-            (DispatchDevice::Autodiff(a), DispatchDevice::Autodiff(b)) => a == b,
+            (DispatchDevice::Autodiff(a), DispatchDevice::Autodiff(b)) => {
+                a.inner.as_ref() == b.inner.as_ref()
+            }
             // If one is Autodiff, compare it to the raw device
             #[cfg(feature = "autodiff")]
             (DispatchDevice::Autodiff(a), b) => a.inner.as_ref() == b,
@@ -421,13 +455,14 @@ impl DispatchDevice {
     #[cfg(feature = "autodiff")]
     /// Creates a new [`DispatchDevice`] with [automatic differentiation](Autodiff) enabled.
     pub fn autodiff(device: impl Into<DispatchDevice>) -> DispatchDevice {
-        Self::autodiff_checkpointed(device, CheckpointingStrategy::None)
+        Self::autodiff_with_gradient_checkpointing(device, GradientCheckpointingStrategy::Disabled)
     }
     #[cfg(feature = "autodiff")]
-    /// Creates a new [`DispatchDevice`] with [automatic differentiation](Autodiff) enabled.
-    pub fn autodiff_checkpointed(
+    /// Creates a new [`DispatchDevice`] with automatic differentiation and the provided gradient
+    /// checkpointing strategy enabled.
+    pub fn autodiff_with_gradient_checkpointing(
         device: impl Into<DispatchDevice>,
-        checkpointing: CheckpointingStrategy,
+        checkpointing: GradientCheckpointingStrategy,
     ) -> DispatchDevice {
         let device = device.into();
         DispatchDevice::Autodiff(AutodiffDevice::new(device, checkpointing))

crates/burn-dispatch/src/macros.rs

@@ -60,14 +60,14 @@ macro_rules! backend_matrix {
 macro_rules! with_autodiff_backend {
     ($Backend:ident, $checkpointing:expr, |$B:ident| $body:expr) => {
         match $checkpointing {
-            Some($crate::CheckpointingStrategy::Balanced) => {
+            Some($crate::GradientCheckpointingStrategy::Balanced) => {
                 type $B = $crate::backends::Autodiff<
                     $crate::backends::$Backend,
                     burn_autodiff::checkpoint::strategy::BalancedCheckpointing,
                 >;
                 $body
             }
-            Some($crate::CheckpointingStrategy::None) => {
+            Some($crate::GradientCheckpointingStrategy::Disabled) => {
                 type $B = $crate::backends::Autodiff<
                     $crate::backends::$Backend,
                     burn_autodiff::checkpoint::strategy::NoCheckpointing,

crates/burn-dispatch/src/tensor.rs

@@ -6,7 +6,7 @@ use burn_autodiff::checkpoint::strategy::{
 };
 use burn_backend::{Backend, BackendTypes, DType, Shape, TensorMetadata};
 
-use crate::CheckpointingStrategy;
+use crate::GradientCheckpointingStrategy;
 #[cfg(feature = "autodiff")]
 use alloc::boxed::Box;
 #[cfg(feature = "autodiff")]
@@ -211,7 +211,7 @@ pub struct DispatchTensor {
     /// Holds the autodiff checkpointing strategy.
     /// - `None`: tensor is not tracked by autodiff
     /// - `Some(strategy)`: tensor is tracked by autodiff, and uses the checkpointing `strategy`
-    pub checkpointing: Option<CheckpointingStrategy>,
+    pub checkpointing: Option<GradientCheckpointingStrategy>,
 }
 
 /// Internal representation of a [`DispatchTensor`].
@@ -477,18 +477,18 @@ impl DispatchTensorKind {
 }
 
 #[cfg(feature = "autodiff")]
-trait IntoCheckpointingStrategy {
-    const STRATEGY: CheckpointingStrategy;
+trait IntoGradientCheckpointingStrategy {
+    const STRATEGY: GradientCheckpointingStrategy;
 }
 
 #[cfg(feature = "autodiff")]
-impl IntoCheckpointingStrategy for NoCheckpointing {
-    const STRATEGY: CheckpointingStrategy = CheckpointingStrategy::None;
+impl IntoGradientCheckpointingStrategy for NoCheckpointing {
+    const STRATEGY: GradientCheckpointingStrategy = GradientCheckpointingStrategy::Disabled;
 }
 
 #[cfg(feature = "autodiff")]
-impl IntoCheckpointingStrategy for BalancedCheckpointing {
-    const STRATEGY: CheckpointingStrategy = CheckpointingStrategy::Balanced;
+impl IntoGradientCheckpointingStrategy for BalancedCheckpointing {
+    const STRATEGY: GradientCheckpointingStrategy = GradientCheckpointingStrategy::Balanced;
 }
 
 /// Trait to execute runtime routing conversions between the dynamic dispatch layer and specific backends.
@@ -531,7 +531,7 @@ macro_rules! impl_dispatch_conversion {
         }
 
         #[cfg(all($cfg, feature = "autodiff"))]
-        impl<C: CheckpointStrategy + IntoCheckpointingStrategy>
+        impl<C: CheckpointStrategy + IntoGradientCheckpointingStrategy>
             DispatchKindConversion<Autodiff<$backend, C>> for DispatchTensor
         {
             fn try_into_backend(

crates/burn-fusion/src/backend.rs

@@ -215,6 +215,11 @@ pub trait NumOperations: core::fmt::Debug {
     }
     /// The name of the optimization.
     fn name(&self) -> &'static str;
+    /// The highest relative shape id this optimization names, if it names any.
+    /// `None` means the optimization names no relative shape, so it fits any stream.
+    fn max_relative_shape_id(&self) -> Option<usize> {
+        None
+    }
 }
 
 /// The optimization created from a [fuser](OperationFuser).

crates/burn-fusion/src/stream/context.rs

@@ -147,6 +147,15 @@ pub(crate) trait RelativeOps {
 }
 
 impl OperationConverter {
+    /// How many relative shape ids have been handed out so far.
+    ///
+    /// Ids are dense from zero, so this is also one past the highest valid id, and therefore
+    /// what a cached plan's
+    /// [`max_relative_shape_id`](crate::NumOperations::max_relative_shape_id) has to fit under.
+    pub(crate) fn num_relative_shapes(&self) -> usize {
+        self.shapes_relative2global.len()
+    }
+
     pub(crate) fn clear(&mut self) {
         self.tensors_relative2global.clear();
         self.tensors_global2relative.clear();

crates/burn-fusion/src/stream/execution/policy.rs

@@ -81,8 +81,12 @@ impl<O: core::fmt::Debug> Policy<O> {
             );
         }
 
-        if let Some((id, _length)) = self.found {
-            return Action::Execute(id);
+        if let Some((id, length)) = self.found {
+            // A plan covering `length` operations cannot be applied to a shorter
+            // segment.
+            if length <= operations.len() {
+                return Action::Execute(id);
+            }
         }
 
         match mode {
@@ -272,6 +276,22 @@ mod tests {
         stream::store::{ExecutionPlan, ExecutionStrategy, ExecutionTrigger},
     };
     use std::ops::Range;
+    use std::sync::Arc;
+
+    #[derive(Debug)]
+    struct ShapeIds(Option<usize>);
+
+    impl crate::NumOperations for ShapeIds {
+        fn len(&self) -> usize {
+            1
+        }
+        fn name(&self) -> &'static str {
+            "shape-ids"
+        }
+        fn max_relative_shape_id(&self) -> Option<usize> {
+            self.0
+        }
+    }
 
     #[test]
     fn given_no_optimization_should_explore() {
@@ -315,6 +335,58 @@ mod tests {
         assert_eq!(action, Action::Execute(id_1));
     }
 
+    #[test]
+    fn should_not_execute_a_plan_longer_than_the_segment() {
+        let mut store = ExecutionPlanStore::<()>::default();
+        let mut policy = Policy::<()>::new();
+        let stream = TestStream::new(3);
+
+        let id = store.add(ExecutionPlan {
+            operations: stream.operations[0..3].to_vec(),
+            triggers: vec![ExecutionTrigger::Always],
+            optimization: BlockOptimization::new(ExecutionStrategy::operations(3), Vec::new()),
+        });
+
+        for operation in stream.operations[0..3].iter() {
+            policy.update(&store, operation);
+        }
+
+        assert_eq!(
+            policy.action(&store, &stream.operations[0..3], ExecutionMode::Lazy),
+            Action::Execute(id),
+        );
+        assert_ne!(
+            policy.action(&store, &stream.operations[0..1], ExecutionMode::Lazy),
+            Action::Execute(id),
+        );
+    }
+
+    #[test]
+    fn strategy_reports_the_highest_relative_shape_id_it_holds() {
+        let strategy = ExecutionStrategy::Composed(vec![
+            Box::new(ExecutionStrategy::operations(2)),
+            Box::new(ExecutionStrategy::Optimization {
+                opt: ShapeIds(Some(4)),
+                ordering: Arc::new(vec![0, 1]),
+                score: 0,
+            }),
+            Box::new(ExecutionStrategy::Optimization {
+                opt: ShapeIds(Some(7)),
+                ordering: Arc::new(vec![2]),
+                score: 0,
+            }),
+        ]);
+
+        assert_eq!(strategy.max_relative_shape_id(), Some(7));
+    }
+
+    #[test]
+    fn operations_only_strategy_names_no_relative_shape_id() {
+        let strategy = ExecutionStrategy::<ShapeIds>::operations(3);
+
+        assert_eq!(strategy.max_relative_shape_id(), None);
+    }
+
     #[test]
     fn given_existing_plan_when_found_trigger_should_execute_plan() {
         let mut store = ExecutionPlanStore::default();

crates/burn-fusion/src/stream/queue/base.rs

@@ -20,6 +20,14 @@ pub struct OperationQueue<R: FusionRuntime> {
     /// because we don't need to know the exact values, but they are sufficient to
     /// determine which operations can be fused.
     pub(crate) relative: Vec<OperationIr>,
+    /// `shapes_assigned[i]` is [`OperationConverter::num_relative_shapes`] right after
+    /// `relative[i]` was relativized, so a plan covering the first `n` operations may only
+    /// name relative shape ids below `shapes_assigned[n - 1]`.
+    ///
+    /// The converter's live count answers a different question — it covers the *whole* queue,
+    /// including operations queued after the ones a plan replays — so it cannot be used to
+    /// bound a plan (see [`execute`](Self::execute)).
+    pub(crate) shapes_assigned: Vec<usize>,
     pub(crate) converter: OperationConverter,
     pub(crate) operations: Vec<UnfusedOp<R>>,
     pub(crate) variables: HashMap<TensorId, TensorStatus>,
@@ -41,6 +49,7 @@ impl<R: FusionRuntime> OperationQueue<R> {
         Self {
             global: Vec::new(),
             relative: Vec::new(),
+            shapes_assigned: Vec::new(),
             converter: OperationConverter::default(),
             operations: Vec::new(),
             variables: HashMap::new(),
@@ -88,6 +97,8 @@ impl<R: FusionRuntime> OperationQueue<R> {
         }
         let relative = global.to_relative(&mut self.converter);
         self.relative.push(relative);
+        self.shapes_assigned
+            .push(self.converter.num_relative_shapes());
         self.global.push(global);
         self.operations.push(operation);
     }

crates/burn-fusion/src/stream/queue/execution.rs

@@ -1,4 +1,6 @@
 use burn_ir::{HandleContainer, TensorStatus};
+use burn_std::config::{fusion::FusionLogLevel, log_fusion};
+use std::sync::Arc;
 
 use crate::{
     FusionRuntime, UnfusedOp,
@@ -22,6 +24,42 @@ impl<R: FusionRuntime> OperationQueue<R> {
         stream_id: StreamId,
     ) {
         let plan = store.get_mut_unchecked(id);
+
+        // A cached plan may name relative shape ids this stream never assigned. Matching
+        // on operations therefore does not imply the plan fits. When it does not, run the very
+        // same operations in submission order instead: always a legal order, just unfused.
+        //
+        // The bound is what the plan's own operations assigned, not what the whole queue did:
+        // plans usually fire from `ExecutionTrigger::OnOperations`, i.e. exactly when later
+        // operations are already queued behind them, and those would otherwise inflate the
+        // count enough to let an unfitting plan through.
+        let len = plan.optimization.ordering.len();
+        let assigned = match len.checked_sub(1).and_then(|i| self.shapes_assigned.get(i)) {
+            Some(assigned) => *assigned,
+            // No operation to bound against: nothing but shape id 0 can be legal.
+            None => 1,
+        };
+        if let Some(max_id) = plan.optimization.strategy.max_relative_shape_id()
+            && max_id >= assigned
+        {
+            log_fusion(FusionLogLevel::Medium, || {
+                format!(
+                    "[plan] #{id} needs relative shape id {max_id} but the stream assigned \
+                     {assigned}; running its {len} operations unfused"
+                )
+            });
+
+            let ordering: Vec<usize> = (0..len).collect();
+            let mut fallback = BlockOptimization::new(
+                ExecutionStrategy::Operations {
+                    ordering: Arc::new(ordering.clone()),
+                },
+                ordering,
+            );
+            self.execute_block_optimization(&mut fallback, handles, stream_id);
+            return;
+        }
+
         self.execute_block_optimization(&mut plan.optimization, handles, stream_id);
     }
 
@@ -77,11 +115,14 @@ impl<R: FusionRuntime> OperationQueue<R> {
 
     fn reset_relative(&mut self) {
         self.relative.clear();
+        self.shapes_assigned.clear();
         self.converter.clear();
 
         for node in self.global.iter() {
             let relative = node.to_relative(&mut self.converter);
             self.relative.push(relative);
+            self.shapes_assigned
+                .push(self.converter.num_relative_shapes());
         }
     }
 }

crates/burn-fusion/src/stream/store/base.rs

@@ -28,6 +28,19 @@ pub(crate) enum ExecutionStrategy<O> {
     Composed(Vec<Box<Self>>),
 }
 
+impl<O: crate::NumOperations> ExecutionStrategy<O> {
+    pub(crate) fn max_relative_shape_id(&self) -> Option<usize> {
+        match self {
+            Self::Optimization { opt, .. } => opt.max_relative_shape_id(),
+            Self::Operations { .. } => None,
+            Self::Composed(items) => items
+                .iter()
+                .filter_map(|item| item.max_relative_shape_id())
+                .max(),
+        }
+    }
+}
+
 /// The trigger that indicates when to stop exploring.
 #[derive(Debug, PartialEq, Serialize, Deserialize)]
 pub(crate) enum ExecutionTrigger {

crates/burn-optim/src/optim/visitor.rs

@@ -34,6 +34,11 @@ where
             return;
         }
 
+        // Prevents `grad_remove` from panicking if one of the tensor is not autodiff.
+        if !param.is_require_grad() {
+            return;
+        }
+
         let Some(grad) = param.val().grad_remove(self.grads) else {
             return;
         };

crates/burn-tensor/src/device.rs

@@ -7,6 +7,8 @@ pub use burn_backend::cubecl::{ThroughputKey, ThroughputMode, ThroughputValue};
 use burn_backend::{Backend, DeviceOps};
 #[allow(unused)]
 use burn_dispatch::DispatchDeviceId;
+#[cfg(feature = "autodiff")]
+use burn_dispatch::GradientCheckpointingStrategy;
 use burn_dispatch::{Dispatch, DispatchDevice};
 use burn_std::{BoolDType, FloatDType, IntDType, TensorData};
 
@@ -538,6 +540,19 @@ impl Device {
         }
     }
 
+    /// Returns an autodiff device's gradient checkpointing strategy.
+    ///
+    /// # Panics
+    ///
+    /// Panics if autodiff is not enabled on this device.
+    #[cfg(feature = "autodiff")]
+    pub fn gradient_checkpointing_strategy(&self) -> GradientCheckpointingStrategy {
+        match self.as_dispatch() {
+            DispatchDevice::Autodiff(device) => device.gradient_checkpointing_strategy(),
+            _ => panic!("Autodiff is not enabled on this device"),
+        }
+    }
+
     /// Enables gradient checkpointing on the autodiff device.
     ///
     /// Gradient checkpointing recomputes activations during backpropagation for operations
@@ -558,11 +573,9 @@ impl Device {
     pub fn gradient_checkpointing(self) -> Self {
         match self.into_dispatch() {
             DispatchDevice::Autodiff(device) => {
-                use burn_dispatch::CheckpointingStrategy;
-
-                Self::new(DispatchDevice::autodiff_checkpointed(
+                Self::new(DispatchDevice::autodiff_with_gradient_checkpointing(
                     device.inner(),
-                    CheckpointingStrategy::Balanced,
+                    GradientCheckpointingStrategy::Balanced,
                 ))
             }
             _ => panic!("Autodiff is not enabled on this device"),

crates/burn-tensor/src/lib.rs

@@ -44,6 +44,8 @@ pub(crate) use tensor::check::macros::check;
 pub use tensor::*;
 
 // Re-exported types
+#[cfg(feature = "autodiff")]
+pub use burn_dispatch::GradientCheckpointingStrategy;
 pub use burn_std::{
     AllocationProperty, Bytes, bf16, f16,
     reader::{read_sync, try_read_sync},

crates/burn-tensor/src/tensor/api/autodiff.rs

@@ -1,11 +1,13 @@
 use crate::{Tensor, kind::Autodiff};
 
 #[cfg(feature = "autodiff")]
-use crate::ops::BridgeTensor;
+use crate::ops::{BridgeKind, BridgeTensor};
 #[cfg(feature = "autodiff")]
 use burn_backend::AutodiffBackend;
 #[cfg(feature = "autodiff")]
 use burn_dispatch::Dispatch;
+#[cfg(feature = "autodiff")]
+use burn_dispatch::GradientCheckpointingStrategy;
 
 #[cfg(feature = "autodiff")]
 type AutodiffGradients = <Dispatch as AutodiffBackend>::Gradients;
@@ -132,6 +134,32 @@ impl<const D: usize, K: Autodiff> Tensor<D, K> {
     pub fn from_inner(inner: Tensor<D, K>) -> Self {
         Self::new(K::from_inner(inner.primitive))
     }
+
+    /// Sets the autodiff checkpointing strategy carried by this tensor.
+    ///
+    /// The strategy is normally derived from the device the tensor was created on (see
+    /// [`Device::gradient_checkpointing`](crate::Device::gradient_checkpointing)); this
+    /// method overrides it for a single tensor. A tensor carrying a strategy is treated
+    /// as tracked by autodiff, so this also marks an inner-backend tensor for tracking.
+    ///
+    /// # Panics
+    ///
+    /// Operations combining tensors that carry different strategies panic; make sure all
+    /// operands share the same one.
+    #[cfg(feature = "autodiff")]
+    pub fn with_gradient_checkpointing_strategy(
+        self,
+        strategy: GradientCheckpointingStrategy,
+    ) -> Self {
+        let (kind, mut tensor) = self.primitive.into_parts();
+        tensor.checkpointing = Some(strategy);
+        Self::new(match kind {
+            BridgeKind::Bool => BridgeTensor::bool(tensor),
+            BridgeKind::Int => BridgeTensor::int(tensor),
+            BridgeKind::Float => BridgeTensor::float(tensor),
+            BridgeKind::QFloat => BridgeTensor::qfloat(tensor),
+        })
+    }
 }
 
 // TODO: a lot of the `tensor.inner` / `Tensor::from_inner(...)` are actually scoped to perform some operations

crates/burn-train/src/learner/rl/paradigm.rs

@@ -150,11 +150,8 @@ impl<RLC: RLComponentsTypes + 'static> RLTraining<RLC> {
     /// # Arguments
     ///
     /// * `renderer` - The custom renderer.
-    pub fn renderer<MR>(mut self, renderer: MR) -> Self
-    where
-        MR: MetricsRenderer + 'static,
-    {
-        self.renderer = Some(Box::new(renderer));
+    pub fn renderer(mut self, renderer: Box<dyn MetricsRenderer + 'static>) -> Self {
+        self.renderer = Some(renderer);
         self
     }
 

crates/burn-train/src/learner/supervised/paradigm.rs

@@ -176,11 +176,8 @@ impl<M: LearnerModel> SupervisedTraining<M> {
     /// # Arguments
     ///
     /// * `renderer` - The custom renderer.
-    pub fn renderer<MR>(mut self, renderer: MR) -> Self
-    where
-        MR: MetricsRenderer + 'static,
-    {
-        self.renderer = Some(Box::new(renderer));
+    pub fn renderer(mut self, renderer: Box<dyn MetricsRenderer + 'static>) -> Self {
+        self.renderer = Some(renderer);
         self
     }
 

crates/burn-train/src/renderer/mod.rs

@@ -20,10 +20,9 @@ use crate::Interrupter;
 ///     a terminal, or
 ///   - `CliMetricsRenderer`, when the `tui` feature is not enabled, or `stdout`
 ///     is not a terminal.
-#[allow(unused_variables)]
-pub(crate) fn default_renderer(
-    interuptor: Interrupter,
-    checkpoint: Option<usize>,
+pub fn default_renderer(
+    #[cfg_attr(not(feature = "tui"), allow(unused_variables))] interuptor: Interrupter,
+    #[cfg_attr(not(feature = "tui"), allow(unused_variables))] checkpoint: Option<usize>,
 ) -> Box<dyn MetricsRenderer> {
     #[cfg(feature = "tui")]
     if std::io::stdout().is_terminal() {

crates/burn-train/src/renderer/tui/base.rs

@@ -33,9 +33,12 @@ impl MetricsView<'_> {
         let size_other = chunks[0];
         let size_metric_numeric = chunks[1];
 
+        // TODO: constraints are still hardcoded, but could be computed dynamically. For example,
+        // the constraints on `size_status` could depend on the number of different progress events
+        // logged with [`TrainingProgressLogger::log_event_training`].
         let chunks = Layout::default()
             .direction(Direction::Vertical)
-            .constraints([Constraint::Max(5), Constraint::Min(6), Constraint::Max(6)].as_ref())
+            .constraints([Constraint::Max(5), Constraint::Min(6), Constraint::Max(7)].as_ref())
             .split(size_other);
         let size_controls = chunks[0];
         let size_metric_text = chunks[1];

crates/burn-vision/src/filter.rs

@@ -105,6 +105,11 @@ impl BoxBlur {
         self
     }
 
+    /// Rebuilds the precomputed kernel on `device`.
+    pub fn to_device(&mut self, device: &Device) {
+        self.kernel = self.kernel.clone().to_device(device);
+    }
+
     fn sample(&self) -> bool {
         let mut rng = self.rng.lock();
 

examples/custom-renderer/src/lib.rs

@@ -121,7 +121,7 @@ pub fn run(device: Device) {
     // artifact dir does not need to be provided when log_to_file is false
     let training = SupervisedTraining::new("", dataloader_train, dataloader_test)
         .num_epochs(config.num_epochs)
-        .renderer(CustomRenderer {})
+        .renderer(Box::new(CustomRenderer {}))
         .with_application_logger(None);
     // can be used to interrupt training
     let _interrupter = training.interrupter();

examples/mnist/cubecl.toml

@@ -1,15 +0,0 @@
-[profiling]
-logger = { log = "info", level = "disabled" }
-
-[autotune]
-level = "balanced"
-cache = "target"
-logger = { file = "/tmp/autotune.log", level = "disabled" }
-
-[compilation]
-logger = { level = "disabled" }
-cache = "target"
-
-[memory]
-logger = { level = "disabled", file = "/tmp/memory.log" }
-persistent_memory = "enabled"