burn_flex/
qtensor.rs

1use alloc::vec::Vec;
2
3use burn_backend::{DType, TensorMetadata};
4use burn_std::{QuantScheme, Shape, quantization::global_scale_dtype};
5
6use crate::{FlexDevice, tensor::FlexTensor};
7
8/// Quantized tensor for the Flex backend.
9///
10/// Stores quantized i8 values in the tensor and keeps scales separately
11/// for efficient dequantization without reparsing bytes.
12#[derive(Clone, Debug)]
13pub struct FlexQTensor {
14    /// The underlying quantized data (stored as i8).
15    pub(crate) tensor: FlexTensor,
16    /// Quantization scheme.
17    pub(crate) scheme: QuantScheme,
18    /// Per-tensor or per-block scale factors.
19    pub(crate) scales: Vec<f32>,
20    /// The per-tensor scale that [`scales`](Self::scales) are expressed relative to, for a
21    /// two-level scheme.
22    pub(crate) global: Option<f32>,
23}
24
25impl FlexQTensor {
26    /// Create a new quantized tensor.
27    ///
28    /// The tensor must store i8 data and scales must be non-empty. A two-level scheme takes a
29    /// per-tensor scale, and no other level does.
30    pub fn new(
31        tensor: FlexTensor,
32        scheme: QuantScheme,
33        scales: Vec<f32>,
34        global: Option<f32>,
35    ) -> Self {
36        assert_eq!(
37            tensor.dtype(),
38            DType::I8,
39            "quantized tensor must store i8 data, got {:?}",
40            tensor.dtype()
41        );
42        assert!(
43            !scales.is_empty(),
44            "quantized tensor must have at least one scale factor"
45        );
46        assert_eq!(
47            global_scale_dtype(&scheme).is_some(),
48            global.is_some(),
49            "{scheme:?} does not match a per-tensor scale of {global:?}"
50        );
51        Self {
52            tensor,
53            scheme,
54            scales,
55            global,
56        }
57    }
58
59    /// Get the underlying tensor.
60    pub fn tensor(&self) -> &FlexTensor {
61        &self.tensor
62    }
63
64    /// Get the quantization scales.
65    pub fn scales(&self) -> &[f32] {
66        &self.scales
67    }
68}
69
70impl TensorMetadata for FlexQTensor {
71    type Device = FlexDevice;
72
73    fn dtype(&self) -> DType {
74        DType::QFloat(self.scheme)
75    }
76
77    fn shape(&self) -> Shape {
78        self.tensor.shape()
79    }
80
81    fn rank(&self) -> usize {
82        self.tensor.rank()
83    }
84
85    fn device(&self) -> Self::Device {
86        FlexDevice
87    }
88
89    fn can_mut(&self) -> bool {
90        self.tensor.is_unique()
91    }
92}
burn_flex/qtensor.rs

burn_flex/
qtensor.rs
