burn_backend/backend/
primitive.rs

1use crate::{Backend, BackendTypes, DeviceOps, get_device_settings};
2use burn_std::{DType, QuantScheme, Shape};
3
4#[derive(Debug, Clone)]
5/// A primitive tensor representation.
6pub enum TensorPrimitive<B: BackendTypes> {
7    /// Float tensor primitive.
8    Float(B::FloatTensorPrimitive),
9    /// Quantized float tensor primitive.
10    QFloat(B::QuantizedTensorPrimitive),
11}
12
13impl<B: Backend> TensorPrimitive<B> {
14    /// Returns the full tensor representation.
15    pub fn tensor(self) -> B::FloatTensorPrimitive {
16        match self {
17            Self::QFloat(tensor) => {
18                let dtype = get_device_settings::<B>(&tensor.device()).float_dtype;
19                B::dequantize(tensor, dtype)
20            }
21            Self::Float(tensor) => tensor,
22        }
23    }
24
25    /// Returns a mutable reference to the full tensor representation.
26    pub fn get_mut_ref(&mut self) -> &mut B::FloatTensorPrimitive {
27        match self {
28            Self::QFloat(_tensor) => todo!(),
29            Self::Float(tensor) => tensor,
30        }
31    }
32}
33
34impl<B: BackendTypes> TensorMetadata for TensorPrimitive<B> {
35    type Device = B::Device;
36
37    fn dtype(&self) -> DType {
38        match self {
39            TensorPrimitive::Float(tensor) => tensor.dtype(),
40            TensorPrimitive::QFloat(tensor) => tensor.dtype(),
41        }
42    }
43
44    fn shape(&self) -> Shape {
45        match self {
46            TensorPrimitive::Float(tensor) => tensor.shape(),
47            TensorPrimitive::QFloat(tensor) => tensor.shape(),
48        }
49    }
50
51    fn rank(&self) -> usize {
52        match self {
53            TensorPrimitive::Float(tensor) => tensor.rank(),
54            TensorPrimitive::QFloat(tensor) => tensor.rank(),
55        }
56    }
57    fn device(&self) -> Self::Device {
58        match self {
59            TensorPrimitive::Float(tensor) => tensor.device(),
60            TensorPrimitive::QFloat(tensor) => tensor.device(),
61        }
62    }
63
64    fn can_mut(&self) -> bool {
65        match self {
66            TensorPrimitive::Float(tensor) => tensor.can_mut(),
67            TensorPrimitive::QFloat(tensor) => tensor.can_mut(),
68        }
69    }
70}
71
72/// Tensor metadata trait for tensor primitive.
73pub trait TensorMetadata: Clone + Send + Sync + core::fmt::Debug {
74    /// The device type associated with the tensor.
75    type Device: DeviceOps;
76    /// Get the dtype of the tensor.
77    fn dtype(&self) -> DType;
78    /// Get the shape of the tensor.
79    fn shape(&self) -> Shape;
80
81    /// Get the number of dimensions of the tensor.
82    fn rank(&self) -> usize {
83        self.shape().num_dims()
84    }
85    /// Get the device associated with the tensor.
86    fn device(&self) -> Self::Device;
87
88    /// Whether the tensor's buffer can be mutated in place — i.e. this handle
89    /// uniquely owns it, so an in-place op (`slice_assign`, an inplace kernel)
90    /// writes the existing allocation instead of copying it first.
91    ///
92    /// Backends that track buffer ownership (cubecl, fusion, tch) answer
93    /// precisely; a backend that can't must return a conservative `false` —
94    /// the buffer may be aliased, so an in-place write can't be assumed safe.
95    fn can_mut(&self) -> bool;
96
97    /// Get the [quantization scheme](QuantScheme) for a quantized float tensor.
98    ///
99    /// # Panics
100    /// Panics if the tensor is not quantized.
101    fn scheme(&self) -> QuantScheme {
102        match self.dtype() {
103            DType::QFloat(scheme) => scheme,
104            other => panic!("Quantization scheme is not valid for dtype {other:?}"),
105        }
106    }
107}
burn_backend/backend/primitive.rs

burn_backend/backend/
primitive.rs
