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 DispatchDevice::Autodiff(device) = &mut device
420            && let Some(checkpointing) = &self.checkpointing
421        {
422            device.checkpointing = *checkpointing;
423        }
424
425        device
426    }
427}
428
429impl DispatchTensorKind {
430    /// Returns the backend tensor kind name.
431    pub(crate) fn name(&self) -> &'static str {
432        match self {
433            #[cfg(feature = "cpu")]
434            DispatchTensorKind::Cpu(_) => "Cpu",
435            #[cfg(feature = "cuda")]
436            DispatchTensorKind::Cuda(_) => "Cuda",
437            #[cfg(feature = "metal")]
438            DispatchTensorKind::Metal(_) => "Metal",
439            #[cfg(feature = "rocm")]
440            DispatchTensorKind::Rocm(_) => "Rocm",
441            #[cfg(feature = "vulkan")]
442            DispatchTensorKind::Vulkan(_) => "Vulkan",
443            #[cfg(feature = "wgpu")]
444            DispatchTensorKind::Wgpu(_) => "Wgpu",
445            #[cfg(feature = "webgpu")]
446            DispatchTensorKind::WebGpu(_) => "WebGpu",
447            #[cfg(any(feature = "flex", default_backend))]
448            DispatchTensorKind::Flex(_) => "Flex",
449            #[cfg(feature = "ndarray")]
450            DispatchTensorKind::NdArray(_) => "NdArray",
451            #[cfg(feature = "tch")]
452            DispatchTensorKind::LibTorch(_) => "LibTorch",
453            #[cfg(feature = "remote")]
454            DispatchTensorKind::Remote(_) => "Remote",
455            #[cfg(feature = "autodiff")]
456            DispatchTensorKind::Autodiff(_) => "Autodiff",
457        }
458    }
459}
460
461#[cfg(feature = "autodiff")]
462trait IntoCheckpointingStrategy {
463    const STRATEGY: CheckpointingStrategy;
464}
465
466#[cfg(feature = "autodiff")]
467impl IntoCheckpointingStrategy for NoCheckpointing {
468    const STRATEGY: CheckpointingStrategy = CheckpointingStrategy::None;
469}
470
471#[cfg(feature = "autodiff")]
472impl IntoCheckpointingStrategy for BalancedCheckpointing {
473    const STRATEGY: CheckpointingStrategy = CheckpointingStrategy::Balanced;
474}
475
476/// Trait to execute runtime routing conversions between the dynamic dispatch layer and specific backends.
477pub trait DispatchKindConversion<B: Backend> {
478    /// Attempts to extract a backend-specific [`BackendTensor`] wrapper from a generic, dynamically-routed [`DispatchTensor`].
479    ///
480    /// # Errors
481    ///
482    /// Returns an error if the dynamic routing state does not match the requested backend `B`.
483    fn try_into_backend(tensor: DispatchTensor) -> Result<BackendTensor<B>, String>;
484
485    /// Encapsulates a backend-specific tensor variant back into a globally routing [`DispatchTensor`].
486    fn from_backend(tensor: BackendTensor<B>) -> DispatchTensor;
487}
488
489macro_rules! impl_dispatch_conversion {
490    ($backend:ident, $cfg:meta) => {
491        #[cfg($cfg)]
492        impl DispatchKindConversion<$backend> for DispatchTensor {
493            fn try_into_backend(tensor: DispatchTensor) -> Result<BackendTensor<$backend>, String> {
494                // The catch-all is unreachable in single-backend builds (the enum then has one
495                // variant), but required when several backend features are enabled.
496                #[allow(unreachable_patterns)]
497                match tensor.kind {
498                    DispatchTensorKind::$backend(t) => Ok(t),
499                    other => Err(format!(
500                        "Expected {} tensor, got variant: {}",
501                        stringify!($backend),
502                        other.name()
503                    )),
504                }
505            }
506
507            fn from_backend(tensor: BackendTensor<$backend>) -> DispatchTensor {
508                DispatchTensor {
509                    kind: DispatchTensorKind::$backend(tensor),
510                    checkpointing: None,
511                }
512            }
513        }
514
515        #[cfg(all($cfg, feature = "autodiff"))]
516        impl<C: CheckpointStrategy + IntoCheckpointingStrategy>
517            DispatchKindConversion<Autodiff<$backend, C>> for DispatchTensor
518        {
519            fn try_into_backend(
520                tensor: DispatchTensor,
521            ) -> Result<BackendTensor<Autodiff<$backend, C>>, String> {
522                match tensor.kind {
523                    DispatchTensorKind::Autodiff(t) => match *t {
524                        DispatchTensorKind::$backend(t) => match t {
525                            // Encode as `BackendTensor::Float` for `Autodiff<B, C>`
526                            BackendTensor::Autodiff(t) => Ok(BackendTensor::Float(t)),
527                            other => Err(format!(
528                                "Expected Autodiff {} float tensor, got Autodiff variant: {}",
529                                stringify!($backend),
530                                other.name()
531                            )),
532                        },
533                        other => Err(format!(
534                            "Expected Autodiff {} tensor, got Autodiff variant: {}",
535                            stringify!($backend),
536                            other.name()
537                        )),
538                    },
539                    other => Err(format!(
540                        "Expected Autodiff tensor, got backend: {}",
541                        other.name()
542                    )),
543                }
544            }
545
546            fn from_backend(tensor: BackendTensor<Autodiff<$backend, C>>) -> DispatchTensor {
547                // Unwrap the Autodiff backend representation back into the inner hardware representation
548                let kind = match tensor {
549                    // Inverse: Wrap the `Float` variant back into the backend's `Autodiff` primitive variant
550                    BackendTensor::Float(t) => {
551                        let ad_tensor = BackendTensor::Autodiff(t);
552                        // Wrap in the concrete backend's dispatch container
553                        let inner_dispatch = DispatchTensorKind::$backend(ad_tensor);
554                        // Re-apply the outer Autodiff dispatch wrapper
555                        DispatchTensorKind::Autodiff(Box::new(inner_dispatch))
556                    }
557
558                    // Pass-throughs for non-differentiable types
559                    BackendTensor::Int(t) => DispatchTensorKind::$backend(BackendTensor::Int(t)),
560                    BackendTensor::Bool(t) => DispatchTensorKind::$backend(BackendTensor::Bool(t)),
561                    BackendTensor::Quantized(t) => {
562                        DispatchTensorKind::$backend(BackendTensor::Quantized(t))
563                    }
564
565                    BackendTensor::Autodiff(_) => {
566                        panic!("Unexpected Autodiff variant provided to `from_backend`",)
567                    }
568                };
569
570                DispatchTensor {
571                    kind,
572                    checkpointing: Some(C::STRATEGY),
573                }
574            }
575        }
576    };
577}
578
579impl_dispatch_conversion!(Flex, any(feature = "flex", default_backend));
580impl_dispatch_conversion!(Cpu, feature = "cpu");
581impl_dispatch_conversion!(Cuda, feature = "cuda");
582impl_dispatch_conversion!(Rocm, feature = "rocm");
583impl_dispatch_conversion!(Remote, feature = "remote");
584impl_dispatch_conversion!(Metal, feature = "metal");
585impl_dispatch_conversion!(Vulkan, feature = "vulkan");
586impl_dispatch_conversion!(Wgpu, feature = "wgpu");
587impl_dispatch_conversion!(WebGpu, feature = "webgpu");
588impl_dispatch_conversion!(NdArray, feature = "ndarray");
589impl_dispatch_conversion!(LibTorch, feature = "tch");
burn_dispatch/tensor.rs

burn_dispatch/
tensor.rs
