burn_dispatch/
tensor.rs

1use crate::{DispatchDevice, backends::*};
2
3#[cfg(feature = "autodiff")]
4use burn_autodiff::checkpoint::strategy::{
5    BalancedCheckpointing, CheckpointStrategy, NoCheckpointing,
6};
7use burn_backend::{Backend, BackendTypes, DType, Shape, TensorMetadata};
8
9use crate::GradientCheckpointingStrategy;
10#[cfg(feature = "autodiff")]
11use alloc::boxed::Box;
12#[cfg(feature = "autodiff")]
13use burn_backend::tensor::FloatTensor;
14
15use alloc::{format, string::String};
16
17// TODO: if we reduce the different associated types for float/int/bool/quantized tensor primitives down to a single
18// `B::TensorPrimitive` we can simplify this.
19
20/// Tensor which points to a backend tensor primitive kind.
21#[derive(Clone, Debug)]
22pub enum BackendTensor<B: BackendTypes> {
23    /// Float tensor handle.
24    Float(B::FloatTensorPrimitive),
25    /// Int tensor handle.
26    Int(B::IntTensorPrimitive),
27    /// Bool tensor handle.
28    Bool(B::BoolTensorPrimitive),
29    /// Quantized tensor handle.
30    Quantized(B::QuantizedTensorPrimitive),
31    #[cfg(feature = "autodiff")]
32    /// Autodiff float tensor handle.
33    Autodiff(FloatTensor<Autodiff<B>>),
34}
35
36impl<B: Backend> BackendTensor<B> {
37    /// Returns the inner float tensor primitive.
38    pub fn float(self) -> B::FloatTensorPrimitive {
39        match self {
40            BackendTensor::Float(tensor) => tensor,
41            BackendTensor::Int(_) => panic!("Should be float, got int"),
42            BackendTensor::Bool(_) => panic!("Should be float, got bool"),
43            BackendTensor::Quantized(_) => panic!("Should be float, got quantized"),
44            #[cfg(feature = "autodiff")]
45            BackendTensor::Autodiff(_) => panic!("Should be float, got autodiff"),
46        }
47    }
48    /// Returns the inner float tensor primitive.
49    pub fn as_float(&self) -> &B::FloatTensorPrimitive {
50        match self {
51            BackendTensor::Float(tensor) => tensor,
52            BackendTensor::Int(_) => panic!("Should be float, got int"),
53            BackendTensor::Bool(_) => panic!("Should be float, got bool"),
54            BackendTensor::Quantized(_) => panic!("Should be float, got quantized"),
55            #[cfg(feature = "autodiff")]
56            BackendTensor::Autodiff(_) => panic!("Should be float, got autodiff"),
57        }
58    }
59
60    /// Returns the inner int tensor primitive.
61    pub fn int(self) -> B::IntTensorPrimitive {
62        match self {
63            BackendTensor::Int(tensor) => tensor,
64            BackendTensor::Float(_) => panic!("Should be int, got float"),
65            BackendTensor::Bool(_) => panic!("Should be int, got bool"),
66            BackendTensor::Quantized(_) => panic!("Should be int, got quantized"),
67            #[cfg(feature = "autodiff")]
68            BackendTensor::Autodiff(_) => panic!("Should be int, got autodiff"),
69        }
70    }
71
72    /// Returns the inner bool tensor primitive.
73    pub fn bool(self) -> B::BoolTensorPrimitive {
74        match self {
75            BackendTensor::Bool(tensor) => tensor,
76            BackendTensor::Float(_) => panic!("Should be bool, got float"),
77            BackendTensor::Int(_) => panic!("Should be bool, got int"),
78            BackendTensor::Quantized(_) => panic!("Should be bool, got quantized"),
79            #[cfg(feature = "autodiff")]
80            BackendTensor::Autodiff(_) => panic!("Should be bool, got autodiff"),
81        }
82    }
83
84    /// Returns the inner quantized tensor primitive.
85    pub fn quantized(self) -> B::QuantizedTensorPrimitive {
86        match self {
87            BackendTensor::Quantized(tensor) => tensor,
88            _ => unreachable!(),
89        }
90    }
91
92    #[cfg(feature = "autodiff")]
93    /// Returns the inner autodiff tensor primitive.
94    pub fn autodiff(self) -> FloatTensor<Autodiff<B>> {
95        match self {
96            BackendTensor::Autodiff(tensor) => tensor,
97            // NOTE: this is the panicking code reached in tensor.rs:74:18:
98            _ => unreachable!(),
99        }
100    }
101
102    #[cfg(feature = "autodiff")]
103    /// Returns the inner autodiff tensor primitive.
104    pub fn as_autodiff(&self) -> &FloatTensor<Autodiff<B>> {
105        match self {
106            BackendTensor::Autodiff(tensor) => tensor,
107            _ => unreachable!(),
108        }
109    }
110
111    #[cfg(feature = "autodiff")]
112    /// Returns the inner autodiff tensor primitive.
113    pub fn autodiff_inner(self) -> B::FloatTensorPrimitive {
114        match self {
115            BackendTensor::Autodiff(tensor) => tensor.primitive,
116            _ => unreachable!(),
117        }
118    }
119
120    /// Lift a handle for backend `B` into the equivalent handle for `Autodiff<B>`.
121    ///
122    /// An already-tracked float (`Autodiff`) becomes the `Float` handle of `Autodiff<B>`; int/bool/
123    /// quantized handles are re-tagged unchanged (those primitives are shared between `B` and
124    /// `Autodiff<B>`). An untracked `Float` handle is invalid here (under autodiff, float tensors
125    /// arrive tracked), so it panics.
126    #[cfg(feature = "autodiff")]
127    pub fn into_autodiff(self) -> BackendTensor<Autodiff<B>> {
128        match self {
129            BackendTensor::Autodiff(tensor) => BackendTensor::Float(tensor),
130            BackendTensor::Int(tensor) => BackendTensor::Int(tensor),
131            BackendTensor::Bool(tensor) => BackendTensor::Bool(tensor),
132            BackendTensor::Quantized(tensor) => BackendTensor::Quantized(tensor),
133            BackendTensor::Float(_) => {
134                unreachable!("an untracked float handle can't be lifted to Autodiff<B>")
135            }
136        }
137    }
138
139    /// Returns the tensor primitive kind name.
140    pub fn name(&self) -> &'static str {
141        match self {
142            BackendTensor::Float(_) => "Float",
143            BackendTensor::Int(_) => "Int",
144            BackendTensor::Bool(_) => "Bool",
145            BackendTensor::Quantized(_) => "Quantized",
146            #[cfg(feature = "autodiff")]
147            BackendTensor::Autodiff(_) => "Autodiff",
148        }
149    }
150}
151
152impl<B: BackendTypes> TensorMetadata for BackendTensor<B> {
153    type Device = B::Device;
154    fn device(&self) -> Self::Device {
155        match self {
156            BackendTensor::Float(tensor) => tensor.device(),
157            BackendTensor::Int(tensor) => tensor.device(),
158            BackendTensor::Bool(tensor) => tensor.device(),
159            BackendTensor::Quantized(tensor) => tensor.device(),
160            #[cfg(feature = "autodiff")]
161            BackendTensor::Autodiff(tensor) => tensor.device(),
162        }
163    }
164    fn dtype(&self) -> DType {
165        match self {
166            BackendTensor::Float(tensor) => tensor.dtype(),
167            BackendTensor::Int(tensor) => tensor.dtype(),
168            BackendTensor::Bool(tensor) => tensor.dtype(),
169            BackendTensor::Quantized(tensor) => tensor.dtype(),
170            #[cfg(feature = "autodiff")]
171            BackendTensor::Autodiff(tensor) => tensor.dtype(),
172        }
173    }
174
175    fn shape(&self) -> Shape {
176        match self {
177            BackendTensor::Float(tensor) => tensor.shape(),
178            BackendTensor::Int(tensor) => tensor.shape(),
179            BackendTensor::Bool(tensor) => tensor.shape(),
180            BackendTensor::Quantized(tensor) => tensor.shape(),
181            #[cfg(feature = "autodiff")]
182            BackendTensor::Autodiff(tensor) => tensor.shape(),
183        }
184    }
185
186    fn can_mut(&self) -> bool {
187        match self {
188            BackendTensor::Float(tensor) => tensor.can_mut(),
189            BackendTensor::Int(tensor) => tensor.can_mut(),
190            BackendTensor::Bool(tensor) => tensor.can_mut(),
191            BackendTensor::Quantized(tensor) => tensor.can_mut(),
192            #[cfg(feature = "autodiff")]
193            BackendTensor::Autodiff(tensor) => tensor.can_mut(),
194        }
195    }
196}
197
198/// A tensor that can dispatch operations to any enabled backend at runtime.
199///
200/// When the `autodiff` feature is enabled, tensors may carry a checkpointing
201/// strategy used to control gradient computation. This is derived from the
202/// device used to create the tensor.
203#[derive(Clone, Debug)]
204pub struct DispatchTensor {
205    /// Tensor kind primitive.
206    pub kind: DispatchTensorKind,
207    // Technically more of a device property, but device is not a dispatch tensor field.
208    // Right now this is the easiest way to preserve the checkpointing strategy because primitives are not consolidated.
209    // Once float/int/bool primitives are consolidated into a single associative type, we could hold that
210    // property for all autodiff tensors.
211    /// Holds the autodiff checkpointing strategy.
212    /// - `None`: tensor is not tracked by autodiff
213    /// - `Some(strategy)`: tensor is tracked by autodiff, and uses the checkpointing `strategy`
214    pub checkpointing: Option<GradientCheckpointingStrategy>,
215}
216
217/// Internal representation of a [`DispatchTensor`].
218///
219/// This enum contains the concrete backend tensor for each enabled backend.
220/// It is not intended to be used directly; instead, it is manipulated by
221/// the dispatch system to route operations to the correct backend.
222///
223/// Each variant corresponds to a specific backend implementation.
224#[derive(Clone, Debug)]
225pub enum DispatchTensorKind {
226    /// The [CPU backend](Cpu) tensor.
227    #[cfg(feature = "cpu")]
228    Cpu(BackendTensor<Cpu>),
229
230    /// The [CUDA backend](Cuda) tensor.
231    #[cfg(feature = "cuda")]
232    Cuda(BackendTensor<Cuda>),
233
234    /// The [Metal backend](Metal) tensor.
235    #[cfg(feature = "metal")]
236    Metal(BackendTensor<Metal>),
237
238    /// The [ROCm backend](Rocm) tensor.
239    #[cfg(feature = "rocm")]
240    Rocm(BackendTensor<Rocm>),
241
242    /// The [Vulkan backend](Vulkan) tensor.
243    #[cfg(feature = "vulkan")]
244    Vulkan(BackendTensor<Vulkan>),
245
246    /// The [Wgpu backend](Wgpu) tensor.
247    #[cfg(feature = "wgpu")]
248    Wgpu(BackendTensor<Wgpu>),
249
250    /// The [WebGPU backend](Wgpu) tensor.
251    #[cfg(feature = "webgpu")]
252    WebGpu(BackendTensor<WebGpu>),
253
254    /// The [Flex backend](Flex) tensor.
255    #[cfg(any(feature = "flex", default_backend))]
256    Flex(BackendTensor<Flex>),
257
258    /// The [NdArray backend](NdArray) tensor.
259    #[cfg(feature = "ndarray")]
260    NdArray(BackendTensor<NdArray>),
261
262    /// The [LibTorch backend](LibTorch) tensor.
263    #[cfg(feature = "tch")]
264    LibTorch(BackendTensor<LibTorch>),
265
266    /// The [Remote backend](Remote) tensor (lives on a remote server).
267    #[cfg(feature = "remote")]
268    Remote(BackendTensor<Remote>),
269    /// A tensor recorded by the capture backend.
270    #[cfg(feature = "capture")]
271    Capture(BackendTensor<Capture>),
272
273    /// The [autodiff enabled backend](Autodiff) tensor.
274    #[cfg(feature = "autodiff")]
275    Autodiff(Box<DispatchTensorKind>),
276}
277
278impl TensorMetadata for DispatchTensorKind {
279    type Device = DispatchDevice;
280
281    fn dtype(&self) -> DType {
282        match self {
283            #[cfg(feature = "cpu")]
284            Self::Cpu(tensor) => tensor.dtype(),
285            #[cfg(feature = "cuda")]
286            Self::Cuda(tensor) => tensor.dtype(),
287            #[cfg(feature = "metal")]
288            Self::Metal(tensor) => tensor.dtype(),
289            #[cfg(feature = "rocm")]
290            Self::Rocm(tensor) => tensor.dtype(),
291            #[cfg(feature = "vulkan")]
292            Self::Vulkan(tensor) => tensor.dtype(),
293            #[cfg(feature = "wgpu")]
294            Self::Wgpu(tensor) => tensor.dtype(),
295            #[cfg(feature = "webgpu")]
296            Self::WebGpu(tensor) => tensor.dtype(),
297            #[cfg(any(feature = "flex", default_backend))]
298            Self::Flex(tensor) => tensor.dtype(),
299            #[cfg(feature = "ndarray")]
300            Self::NdArray(tensor) => tensor.dtype(),
301            #[cfg(feature = "tch")]
302            Self::LibTorch(tensor) => tensor.dtype(),
303            #[cfg(feature = "remote")]
304            Self::Remote(tensor) => tensor.dtype(),
305            #[cfg(feature = "capture")]
306            Self::Capture(tensor) => tensor.dtype(),
307            #[cfg(feature = "autodiff")]
308            Self::Autodiff(tensor) => tensor.dtype(),
309        }
310    }
311
312    fn shape(&self) -> Shape {
313        match self {
314            #[cfg(feature = "cpu")]
315            Self::Cpu(tensor) => tensor.shape(),
316            #[cfg(feature = "cuda")]
317            Self::Cuda(tensor) => tensor.shape(),
318            #[cfg(feature = "metal")]
319            Self::Metal(tensor) => tensor.shape(),
320            #[cfg(feature = "rocm")]
321            Self::Rocm(tensor) => tensor.shape(),
322            #[cfg(feature = "vulkan")]
323            Self::Vulkan(tensor) => tensor.shape(),
324            #[cfg(feature = "wgpu")]
325            Self::Wgpu(tensor) => tensor.shape(),
326            #[cfg(feature = "webgpu")]
327            Self::WebGpu(tensor) => tensor.shape(),
328            #[cfg(any(feature = "flex", default_backend))]
329            Self::Flex(tensor) => tensor.shape(),
330            #[cfg(feature = "ndarray")]
331            Self::NdArray(tensor) => tensor.shape(),
332            #[cfg(feature = "tch")]
333            Self::LibTorch(tensor) => tensor.shape(),
334            #[cfg(feature = "remote")]
335            Self::Remote(tensor) => tensor.shape(),
336            #[cfg(feature = "capture")]
337            Self::Capture(tensor) => tensor.shape(),
338            #[cfg(feature = "autodiff")]
339            Self::Autodiff(tensor) => tensor.shape(),
340        }
341    }
342
343    fn device(&self) -> DispatchDevice {
344        match self {
345            #[cfg(feature = "cpu")]
346            DispatchTensorKind::Cpu(tensor) => DispatchDevice::Cpu(tensor.device()),
347            #[cfg(feature = "cuda")]
348            DispatchTensorKind::Cuda(tensor) => DispatchDevice::Cuda(tensor.device()),
349            #[cfg(feature = "metal")]
350            DispatchTensorKind::Metal(tensor) => DispatchDevice::Metal(tensor.device()),
351            #[cfg(feature = "rocm")]
352            DispatchTensorKind::Rocm(tensor) => DispatchDevice::Rocm(tensor.device()),
353            #[cfg(feature = "vulkan")]
354            DispatchTensorKind::Vulkan(tensor) => DispatchDevice::Vulkan(tensor.device()),
355            #[cfg(feature = "wgpu")]
356            DispatchTensorKind::Wgpu(tensor) => DispatchDevice::Wgpu(tensor.device()),
357            #[cfg(feature = "webgpu")]
358            DispatchTensorKind::WebGpu(tensor) => DispatchDevice::WebGpu(tensor.device()),
359            #[cfg(any(feature = "flex", default_backend))]
360            DispatchTensorKind::Flex(tensor) => DispatchDevice::Flex(tensor.device()),
361            #[cfg(feature = "ndarray")]
362            DispatchTensorKind::NdArray(tensor) => DispatchDevice::NdArray(tensor.device()),
363            #[cfg(feature = "tch")]
364            DispatchTensorKind::LibTorch(tensor) => DispatchDevice::LibTorch(tensor.device()),
365            #[cfg(feature = "remote")]
366            DispatchTensorKind::Remote(tensor) => DispatchDevice::Remote(tensor.device()),
367            #[cfg(feature = "capture")]
368            DispatchTensorKind::Capture(tensor) => DispatchDevice::Capture(tensor.device()),
369            #[cfg(feature = "autodiff")]
370            DispatchTensorKind::Autodiff(tensor) => DispatchDevice::autodiff(tensor.device()),
371        }
372    }
373
374    fn can_mut(&self) -> bool {
375        match self {
376            #[cfg(feature = "cpu")]
377            Self::Cpu(tensor) => tensor.can_mut(),
378            #[cfg(feature = "cuda")]
379            Self::Cuda(tensor) => tensor.can_mut(),
380            #[cfg(feature = "metal")]
381            Self::Metal(tensor) => tensor.can_mut(),
382            #[cfg(feature = "rocm")]
383            Self::Rocm(tensor) => tensor.can_mut(),
384            #[cfg(feature = "vulkan")]
385            Self::Vulkan(tensor) => tensor.can_mut(),
386            #[cfg(feature = "wgpu")]
387            Self::Wgpu(tensor) => tensor.can_mut(),
388            #[cfg(feature = "webgpu")]
389            Self::WebGpu(tensor) => tensor.can_mut(),
390            #[cfg(any(feature = "flex", default_backend))]
391            Self::Flex(tensor) => tensor.can_mut(),
392            #[cfg(feature = "ndarray")]
393            Self::NdArray(tensor) => tensor.can_mut(),
394            #[cfg(feature = "tch")]
395            Self::LibTorch(tensor) => tensor.can_mut(),
396            #[cfg(feature = "remote")]
397            Self::Remote(tensor) => tensor.can_mut(),
398            #[cfg(feature = "capture")]
399            Self::Capture(tensor) => tensor.can_mut(),
400            #[cfg(feature = "autodiff")]
401            Self::Autodiff(tensor) => tensor.can_mut(),
402        }
403    }
404}
405
406impl TensorMetadata for DispatchTensor {
407    fn dtype(&self) -> DType {
408        self.kind.dtype()
409    }
410
411    fn shape(&self) -> Shape {
412        self.kind.shape()
413    }
414
415    fn can_mut(&self) -> bool {
416        self.kind.can_mut()
417    }
418
419    type Device = DispatchDevice;
420
421    fn device(&self) -> Self::Device {
422        #[allow(unused_mut)]
423        let mut device = self.kind.device();
424
425        #[cfg(feature = "autodiff")]
426        if let Some(checkpointing) = &self.checkpointing {
427            // Int, bool, and quantized tensors travel beside the tape untracked, but
428            // their device must retain the autodiff capability so tensors derived
429            // from them can join the graph (for example, an integer one-hot tensor
430            // cast to float). Plain float gradients can also carry checkpointing
431            // metadata copied from their source tensor, but remain on the inner
432            // backend and must continue to report that device.
433            if !self.dtype().is_float() && !matches!(device, DispatchDevice::Autodiff(_)) {
434                device = DispatchDevice::autodiff(device);
435            }
436            if let DispatchDevice::Autodiff(device) = &mut device {
437                device.checkpointing = *checkpointing;
438            }
439        }
440
441        device
442    }
443}
444
445impl DispatchTensorKind {
446    /// Returns the backend tensor kind name.
447    pub(crate) fn name(&self) -> &'static str {
448        match self {
449            #[cfg(feature = "cpu")]
450            DispatchTensorKind::Cpu(_) => "Cpu",
451            #[cfg(feature = "cuda")]
452            DispatchTensorKind::Cuda(_) => "Cuda",
453            #[cfg(feature = "metal")]
454            DispatchTensorKind::Metal(_) => "Metal",
455            #[cfg(feature = "rocm")]
456            DispatchTensorKind::Rocm(_) => "Rocm",
457            #[cfg(feature = "vulkan")]
458            DispatchTensorKind::Vulkan(_) => "Vulkan",
459            #[cfg(feature = "wgpu")]
460            DispatchTensorKind::Wgpu(_) => "Wgpu",
461            #[cfg(feature = "webgpu")]
462            DispatchTensorKind::WebGpu(_) => "WebGpu",
463            #[cfg(any(feature = "flex", default_backend))]
464            DispatchTensorKind::Flex(_) => "Flex",
465            #[cfg(feature = "ndarray")]
466            DispatchTensorKind::NdArray(_) => "NdArray",
467            #[cfg(feature = "tch")]
468            DispatchTensorKind::LibTorch(_) => "LibTorch",
469            #[cfg(feature = "remote")]
470            DispatchTensorKind::Remote(_) => "Remote",
471            #[cfg(feature = "capture")]
472            DispatchTensorKind::Capture(_) => "Capture",
473            #[cfg(feature = "autodiff")]
474            DispatchTensorKind::Autodiff(_) => "Autodiff",
475        }
476    }
477}
478
479#[cfg(feature = "autodiff")]
480trait IntoGradientCheckpointingStrategy {
481    const STRATEGY: GradientCheckpointingStrategy;
482}
483
484#[cfg(feature = "autodiff")]
485impl IntoGradientCheckpointingStrategy for NoCheckpointing {
486    const STRATEGY: GradientCheckpointingStrategy = GradientCheckpointingStrategy::Disabled;
487}
488
489#[cfg(feature = "autodiff")]
490impl IntoGradientCheckpointingStrategy for BalancedCheckpointing {
491    const STRATEGY: GradientCheckpointingStrategy = GradientCheckpointingStrategy::Balanced;
492}
493
494/// Trait to execute runtime routing conversions between the dynamic dispatch layer and specific backends.
495pub trait DispatchKindConversion<B: Backend> {
496    /// Attempts to extract a backend-specific [`BackendTensor`] wrapper from a generic, dynamically-routed [`DispatchTensor`].
497    ///
498    /// # Errors
499    ///
500    /// Returns an error if the dynamic routing state does not match the requested backend `B`.
501    fn try_into_backend(tensor: DispatchTensor) -> Result<BackendTensor<B>, String>;
502
503    /// Encapsulates a backend-specific tensor variant back into a globally routing [`DispatchTensor`].
504    fn from_backend(tensor: BackendTensor<B>) -> DispatchTensor;
505}
506
507macro_rules! impl_dispatch_conversion {
508    ($backend:ident, $cfg:meta) => {
509        #[cfg($cfg)]
510        impl DispatchKindConversion<$backend> for DispatchTensor {
511            fn try_into_backend(tensor: DispatchTensor) -> Result<BackendTensor<$backend>, String> {
512                // The catch-all is unreachable in single-backend builds (the enum then has one
513                // variant), but required when several backend features are enabled.
514                #[allow(unreachable_patterns)]
515                match tensor.kind {
516                    DispatchTensorKind::$backend(t) => Ok(t),
517                    other => Err(format!(
518                        "Expected {} tensor, got variant: {}",
519                        stringify!($backend),
520                        other.name()
521                    )),
522                }
523            }
524
525            fn from_backend(tensor: BackendTensor<$backend>) -> DispatchTensor {
526                DispatchTensor {
527                    kind: DispatchTensorKind::$backend(tensor),
528                    checkpointing: None,
529                }
530            }
531        }
532
533        #[cfg(all($cfg, feature = "autodiff"))]
534        impl<C: CheckpointStrategy + IntoGradientCheckpointingStrategy>
535            DispatchKindConversion<Autodiff<$backend, C>> for DispatchTensor
536        {
537            fn try_into_backend(
538                tensor: DispatchTensor,
539            ) -> Result<BackendTensor<Autodiff<$backend, C>>, String> {
540                match tensor.kind {
541                    DispatchTensorKind::Autodiff(t) => match *t {
542                        DispatchTensorKind::$backend(t) => match t {
543                            // Encode as `BackendTensor::Float` for `Autodiff<B, C>`
544                            BackendTensor::Autodiff(t) => Ok(BackendTensor::Float(t)),
545                            other => Err(format!(
546                                "Expected Autodiff {} float tensor, got Autodiff variant: {}",
547                                stringify!($backend),
548                                other.name()
549                            )),
550                        },
551                        other => Err(format!(
552                            "Expected Autodiff {} tensor, got Autodiff variant: {}",
553                            stringify!($backend),
554                            other.name()
555                        )),
556                    },
557                    other => Err(format!(
558                        "Expected Autodiff tensor, got backend: {}",
559                        other.name()
560                    )),
561                }
562            }
563
564            fn from_backend(tensor: BackendTensor<Autodiff<$backend, C>>) -> DispatchTensor {
565                // Unwrap the Autodiff backend representation back into the inner hardware representation
566                let kind = match tensor {
567                    // Inverse: Wrap the `Float` variant back into the backend's `Autodiff` primitive variant
568                    BackendTensor::Float(t) => {
569                        let ad_tensor = BackendTensor::Autodiff(t);
570                        // Wrap in the concrete backend's dispatch container
571                        let inner_dispatch = DispatchTensorKind::$backend(ad_tensor);
572                        // Re-apply the outer Autodiff dispatch wrapper
573                        DispatchTensorKind::Autodiff(Box::new(inner_dispatch))
574                    }
575
576                    // Pass-throughs for non-differentiable types
577                    BackendTensor::Int(t) => DispatchTensorKind::$backend(BackendTensor::Int(t)),
578                    BackendTensor::Bool(t) => DispatchTensorKind::$backend(BackendTensor::Bool(t)),
579                    BackendTensor::Quantized(t) => {
580                        DispatchTensorKind::$backend(BackendTensor::Quantized(t))
581                    }
582
583                    BackendTensor::Autodiff(_) => {
584                        panic!("Unexpected Autodiff variant provided to `from_backend`",)
585                    }
586                };
587
588                DispatchTensor {
589                    kind,
590                    checkpointing: Some(C::STRATEGY),
591                }
592            }
593        }
594    };
595}
596
597impl_dispatch_conversion!(Flex, any(feature = "flex", default_backend));
598impl_dispatch_conversion!(Cpu, feature = "cpu");
599impl_dispatch_conversion!(Cuda, feature = "cuda");
600impl_dispatch_conversion!(Rocm, feature = "rocm");
601impl_dispatch_conversion!(Remote, feature = "remote");
602impl_dispatch_conversion!(Capture, feature = "capture");
603impl_dispatch_conversion!(Metal, feature = "metal");
604impl_dispatch_conversion!(Vulkan, feature = "vulkan");
605impl_dispatch_conversion!(Wgpu, feature = "wgpu");
606impl_dispatch_conversion!(WebGpu, feature = "webgpu");
607impl_dispatch_conversion!(NdArray, feature = "ndarray");
608impl_dispatch_conversion!(LibTorch, feature = "tch");
burn_dispatch/tensor.rs

burn_dispatch/
tensor.rs
