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::CheckpointingStrategy;
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<CheckpointingStrategy>,
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
270    /// The [autodiff enabled backend](Autodiff) tensor.
271    #[cfg(feature = "autodiff")]
272    Autodiff(Box<DispatchTensorKind>),
273}
274
275impl TensorMetadata for DispatchTensorKind {
276    type Device = DispatchDevice;
277
278    fn dtype(&self) -> DType {
279        match self {
280            #[cfg(feature = "cpu")]
281            Self::Cpu(tensor) => tensor.dtype(),
282            #[cfg(feature = "cuda")]
283            Self::Cuda(tensor) => tensor.dtype(),
284            #[cfg(feature = "metal")]
285            Self::Metal(tensor) => tensor.dtype(),
286            #[cfg(feature = "rocm")]
287            Self::Rocm(tensor) => tensor.dtype(),
288            #[cfg(feature = "vulkan")]
289            Self::Vulkan(tensor) => tensor.dtype(),
290            #[cfg(feature = "wgpu")]
291            Self::Wgpu(tensor) => tensor.dtype(),
292            #[cfg(feature = "webgpu")]
293            Self::WebGpu(tensor) => tensor.dtype(),
294            #[cfg(any(feature = "flex", default_backend))]
295            Self::Flex(tensor) => tensor.dtype(),
296            #[cfg(feature = "ndarray")]
297            Self::NdArray(tensor) => tensor.dtype(),
298            #[cfg(feature = "tch")]
299            Self::LibTorch(tensor) => tensor.dtype(),
300            #[cfg(feature = "remote")]
301            Self::Remote(tensor) => tensor.dtype(),
302            #[cfg(feature = "autodiff")]
303            Self::Autodiff(tensor) => tensor.dtype(),
304        }
305    }
306
307    fn shape(&self) -> Shape {
308        match self {
309            #[cfg(feature = "cpu")]
310            Self::Cpu(tensor) => tensor.shape(),
311            #[cfg(feature = "cuda")]
312            Self::Cuda(tensor) => tensor.shape(),
313            #[cfg(feature = "metal")]
314            Self::Metal(tensor) => tensor.shape(),
315            #[cfg(feature = "rocm")]
316            Self::Rocm(tensor) => tensor.shape(),
317            #[cfg(feature = "vulkan")]
318            Self::Vulkan(tensor) => tensor.shape(),
319            #[cfg(feature = "wgpu")]
320            Self::Wgpu(tensor) => tensor.shape(),
321            #[cfg(feature = "webgpu")]
322            Self::WebGpu(tensor) => tensor.shape(),
323            #[cfg(any(feature = "flex", default_backend))]
324            Self::Flex(tensor) => tensor.shape(),
325            #[cfg(feature = "ndarray")]
326            Self::NdArray(tensor) => tensor.shape(),
327            #[cfg(feature = "tch")]
328            Self::LibTorch(tensor) => tensor.shape(),
329            #[cfg(feature = "remote")]
330            Self::Remote(tensor) => tensor.shape(),
331            #[cfg(feature = "autodiff")]
332            Self::Autodiff(tensor) => tensor.shape(),
333        }
334    }
335
336    fn device(&self) -> DispatchDevice {
337        match self {
338            #[cfg(feature = "cpu")]
339            DispatchTensorKind::Cpu(tensor) => DispatchDevice::Cpu(tensor.device()),
340            #[cfg(feature = "cuda")]
341            DispatchTensorKind::Cuda(tensor) => DispatchDevice::Cuda(tensor.device()),
342            #[cfg(feature = "metal")]
343            DispatchTensorKind::Metal(tensor) => DispatchDevice::Metal(tensor.device()),
344            #[cfg(feature = "rocm")]
345            DispatchTensorKind::Rocm(tensor) => DispatchDevice::Rocm(tensor.device()),
346            #[cfg(feature = "vulkan")]
347            DispatchTensorKind::Vulkan(tensor) => DispatchDevice::Vulkan(tensor.device()),
348            #[cfg(feature = "wgpu")]
349            DispatchTensorKind::Wgpu(tensor) => DispatchDevice::Wgpu(tensor.device()),
350            #[cfg(feature = "webgpu")]
351            DispatchTensorKind::WebGpu(tensor) => DispatchDevice::WebGpu(tensor.device()),
352            #[cfg(any(feature = "flex", default_backend))]
353            DispatchTensorKind::Flex(tensor) => DispatchDevice::Flex(tensor.device()),
354            #[cfg(feature = "ndarray")]
355            DispatchTensorKind::NdArray(tensor) => DispatchDevice::NdArray(tensor.device()),
356            #[cfg(feature = "tch")]
357            DispatchTensorKind::LibTorch(tensor) => DispatchDevice::LibTorch(tensor.device()),
358            #[cfg(feature = "remote")]
359            DispatchTensorKind::Remote(tensor) => DispatchDevice::Remote(tensor.device()),
360            #[cfg(feature = "autodiff")]
361            DispatchTensorKind::Autodiff(tensor) => DispatchDevice::autodiff(tensor.device()),
362        }
363    }
364
365    fn can_mut(&self) -> bool {
366        match self {
367            #[cfg(feature = "cpu")]
368            Self::Cpu(tensor) => tensor.can_mut(),
369            #[cfg(feature = "cuda")]
370            Self::Cuda(tensor) => tensor.can_mut(),
371            #[cfg(feature = "metal")]
372            Self::Metal(tensor) => tensor.can_mut(),
373            #[cfg(feature = "rocm")]
374            Self::Rocm(tensor) => tensor.can_mut(),
375            #[cfg(feature = "vulkan")]
376            Self::Vulkan(tensor) => tensor.can_mut(),
377            #[cfg(feature = "wgpu")]
378            Self::Wgpu(tensor) => tensor.can_mut(),
379            #[cfg(feature = "webgpu")]
380            Self::WebGpu(tensor) => tensor.can_mut(),
381            #[cfg(any(feature = "flex", default_backend))]
382            Self::Flex(tensor) => tensor.can_mut(),
383            #[cfg(feature = "ndarray")]
384            Self::NdArray(tensor) => tensor.can_mut(),
385            #[cfg(feature = "tch")]
386            Self::LibTorch(tensor) => tensor.can_mut(),
387            #[cfg(feature = "remote")]
388            Self::Remote(tensor) => tensor.can_mut(),
389            #[cfg(feature = "autodiff")]
390            Self::Autodiff(tensor) => tensor.can_mut(),
391        }
392    }
393}
394
395impl TensorMetadata for DispatchTensor {
396    fn dtype(&self) -> DType {
397        self.kind.dtype()
398    }
399
400    fn shape(&self) -> Shape {
401        self.kind.shape()
402    }
403
404    fn can_mut(&self) -> bool {
405        self.kind.can_mut()
406    }
407
408    type Device = DispatchDevice;
409
410    fn device(&self) -> Self::Device {
411        #[allow(unused_mut)]
412        let mut device = self.kind.device();
413
414        // TODO: should int and bool kinds return an autodiff device?
415        // It would be much easier once there is a single underlying primitive type, which
416        // we can wrap with Autodiff in all cases.
417
418        #[cfg(feature = "autodiff")]
419        if let Some(checkpointing) = &self.checkpointing {
420            // A packed (quantized) tensor is float-kind data that travels beside
421            // the tape untracked, so its device is the autodiff one: whatever is
422            // built from it — its dequantized form, LoRA factors over it — must
423            // land on the tape.
424            if matches!(self.dtype(), DType::QFloat(_))
425                && !matches!(device, DispatchDevice::Autodiff(_))
426            {
427                device = DispatchDevice::autodiff(device);
428            }
429            if let DispatchDevice::Autodiff(device) = &mut device {
430                device.checkpointing = *checkpointing;
431            }
432        }
433
434        device
435    }
436}
437
438impl DispatchTensorKind {
439    /// Returns the backend tensor kind name.
440    pub(crate) fn name(&self) -> &'static str {
441        match self {
442            #[cfg(feature = "cpu")]
443            DispatchTensorKind::Cpu(_) => "Cpu",
444            #[cfg(feature = "cuda")]
445            DispatchTensorKind::Cuda(_) => "Cuda",
446            #[cfg(feature = "metal")]
447            DispatchTensorKind::Metal(_) => "Metal",
448            #[cfg(feature = "rocm")]
449            DispatchTensorKind::Rocm(_) => "Rocm",
450            #[cfg(feature = "vulkan")]
451            DispatchTensorKind::Vulkan(_) => "Vulkan",
452            #[cfg(feature = "wgpu")]
453            DispatchTensorKind::Wgpu(_) => "Wgpu",
454            #[cfg(feature = "webgpu")]
455            DispatchTensorKind::WebGpu(_) => "WebGpu",
456            #[cfg(any(feature = "flex", default_backend))]
457            DispatchTensorKind::Flex(_) => "Flex",
458            #[cfg(feature = "ndarray")]
459            DispatchTensorKind::NdArray(_) => "NdArray",
460            #[cfg(feature = "tch")]
461            DispatchTensorKind::LibTorch(_) => "LibTorch",
462            #[cfg(feature = "remote")]
463            DispatchTensorKind::Remote(_) => "Remote",
464            #[cfg(feature = "autodiff")]
465            DispatchTensorKind::Autodiff(_) => "Autodiff",
466        }
467    }
468}
469
470#[cfg(feature = "autodiff")]
471trait IntoCheckpointingStrategy {
472    const STRATEGY: CheckpointingStrategy;
473}
474
475#[cfg(feature = "autodiff")]
476impl IntoCheckpointingStrategy for NoCheckpointing {
477    const STRATEGY: CheckpointingStrategy = CheckpointingStrategy::None;
478}
479
480#[cfg(feature = "autodiff")]
481impl IntoCheckpointingStrategy for BalancedCheckpointing {
482    const STRATEGY: CheckpointingStrategy = CheckpointingStrategy::Balanced;
483}
484
485/// Trait to execute runtime routing conversions between the dynamic dispatch layer and specific backends.
486pub trait DispatchKindConversion<B: Backend> {
487    /// Attempts to extract a backend-specific [`BackendTensor`] wrapper from a generic, dynamically-routed [`DispatchTensor`].
488    ///
489    /// # Errors
490    ///
491    /// Returns an error if the dynamic routing state does not match the requested backend `B`.
492    fn try_into_backend(tensor: DispatchTensor) -> Result<BackendTensor<B>, String>;
493
494    /// Encapsulates a backend-specific tensor variant back into a globally routing [`DispatchTensor`].
495    fn from_backend(tensor: BackendTensor<B>) -> DispatchTensor;
496}
497
498macro_rules! impl_dispatch_conversion {
499    ($backend:ident, $cfg:meta) => {
500        #[cfg($cfg)]
501        impl DispatchKindConversion<$backend> for DispatchTensor {
502            fn try_into_backend(tensor: DispatchTensor) -> Result<BackendTensor<$backend>, String> {
503                // The catch-all is unreachable in single-backend builds (the enum then has one
504                // variant), but required when several backend features are enabled.
505                #[allow(unreachable_patterns)]
506                match tensor.kind {
507                    DispatchTensorKind::$backend(t) => Ok(t),
508                    other => Err(format!(
509                        "Expected {} tensor, got variant: {}",
510                        stringify!($backend),
511                        other.name()
512                    )),
513                }
514            }
515
516            fn from_backend(tensor: BackendTensor<$backend>) -> DispatchTensor {
517                DispatchTensor {
518                    kind: DispatchTensorKind::$backend(tensor),
519                    checkpointing: None,
520                }
521            }
522        }
523
524        #[cfg(all($cfg, feature = "autodiff"))]
525        impl<C: CheckpointStrategy + IntoCheckpointingStrategy>
526            DispatchKindConversion<Autodiff<$backend, C>> for DispatchTensor
527        {
528            fn try_into_backend(
529                tensor: DispatchTensor,
530            ) -> Result<BackendTensor<Autodiff<$backend, C>>, String> {
531                match tensor.kind {
532                    DispatchTensorKind::Autodiff(t) => match *t {
533                        DispatchTensorKind::$backend(t) => match t {
534                            // Encode as `BackendTensor::Float` for `Autodiff<B, C>`
535                            BackendTensor::Autodiff(t) => Ok(BackendTensor::Float(t)),
536                            other => Err(format!(
537                                "Expected Autodiff {} float tensor, got Autodiff variant: {}",
538                                stringify!($backend),
539                                other.name()
540                            )),
541                        },
542                        other => Err(format!(
543                            "Expected Autodiff {} tensor, got Autodiff variant: {}",
544                            stringify!($backend),
545                            other.name()
546                        )),
547                    },
548                    other => Err(format!(
549                        "Expected Autodiff tensor, got backend: {}",
550                        other.name()
551                    )),
552                }
553            }
554
555            fn from_backend(tensor: BackendTensor<Autodiff<$backend, C>>) -> DispatchTensor {
556                // Unwrap the Autodiff backend representation back into the inner hardware representation
557                let kind = match tensor {
558                    // Inverse: Wrap the `Float` variant back into the backend's `Autodiff` primitive variant
559                    BackendTensor::Float(t) => {
560                        let ad_tensor = BackendTensor::Autodiff(t);
561                        // Wrap in the concrete backend's dispatch container
562                        let inner_dispatch = DispatchTensorKind::$backend(ad_tensor);
563                        // Re-apply the outer Autodiff dispatch wrapper
564                        DispatchTensorKind::Autodiff(Box::new(inner_dispatch))
565                    }
566
567                    // Pass-throughs for non-differentiable types
568                    BackendTensor::Int(t) => DispatchTensorKind::$backend(BackendTensor::Int(t)),
569                    BackendTensor::Bool(t) => DispatchTensorKind::$backend(BackendTensor::Bool(t)),
570                    BackendTensor::Quantized(t) => {
571                        DispatchTensorKind::$backend(BackendTensor::Quantized(t))
572                    }
573
574                    BackendTensor::Autodiff(_) => {
575                        panic!("Unexpected Autodiff variant provided to `from_backend`",)
576                    }
577                };
578
579                DispatchTensor {
580                    kind,
581                    checkpointing: Some(C::STRATEGY),
582                }
583            }
584        }
585    };
586}
587
588impl_dispatch_conversion!(Flex, any(feature = "flex", default_backend));
589impl_dispatch_conversion!(Cpu, feature = "cpu");
590impl_dispatch_conversion!(Cuda, feature = "cuda");
591impl_dispatch_conversion!(Rocm, feature = "rocm");
592impl_dispatch_conversion!(Remote, feature = "remote");
593impl_dispatch_conversion!(Metal, feature = "metal");
594impl_dispatch_conversion!(Vulkan, feature = "vulkan");
595impl_dispatch_conversion!(Wgpu, feature = "wgpu");
596impl_dispatch_conversion!(WebGpu, feature = "webgpu");
597impl_dispatch_conversion!(NdArray, feature = "ndarray");
598impl_dispatch_conversion!(LibTorch, feature = "tch");
burn_dispatch/tensor.rs

burn_dispatch/
tensor.rs
