Skip to main content

entrenar/quant/granularity/
params.rs

1//! Quantization parameters and quantized tensor structures
2
3use serde::{Deserialize, Serialize};
4
5use super::{QuantGranularity, QuantMode};
6
7/// Quantization parameters for a tensor
8#[derive(Clone, Debug, Serialize, Deserialize)]
9pub struct QuantParams {
10    /// Scale factor(s)
11    pub scales: Vec<f32>,
12    /// Zero point(s) - empty for symmetric quantization
13    pub zero_points: Vec<i32>,
14    /// Quantization granularity
15    pub granularity: QuantGranularity,
16    /// Quantization mode
17    pub mode: QuantMode,
18    /// Bit width (4 or 8)
19    pub bits: u8,
20}
21
22impl QuantParams {
23    /// Get number of scale/zero-point groups
24    pub fn num_groups(&self) -> usize {
25        self.scales.len()
26    }
27
28    /// Check if asymmetric quantization
29    pub fn is_asymmetric(&self) -> bool {
30        self.mode == QuantMode::Asymmetric
31    }
32}
33
34/// Quantized tensor with per-channel or per-tensor quantization
35#[derive(Clone, Debug, Serialize, Deserialize)]
36pub struct QuantizedTensor {
37    /// Quantized integer data (8-bit representation)
38    pub data: Vec<i8>,
39    /// Quantization parameters
40    pub params: QuantParams,
41    /// Original shape
42    pub shape: Vec<usize>,
43}
44
45impl QuantizedTensor {
46    /// Memory usage in bytes
47    pub fn memory_bytes(&self) -> usize {
48        let data_bytes = self.data.len();
49        let scale_bytes = self.params.scales.len() * 4;
50        let zp_bytes = self.params.zero_points.len() * 4;
51        data_bytes + scale_bytes + zp_bytes
52    }
53}