quantize/methods/symmetric/
mod.rs

1//! Symmetric quantization: one scale per group, codes centered on zero.
2
3use crate::error::{check_bits, check_block, Result};
4use crate::kernels::quantize_sym_packed;
5use crate::packed::Packed;
6use crate::scale::Scale;
7use crate::tensor::Quantized;
8
9/// Quantize `values` into fixed-size blocks of `BLOCK` with `BITS`-wide codes.
10///
11/// # Errors
12///
13/// [`crate::Error::InvalidBits`] or [`crate::Error::InvalidBlock`].
14pub fn quantize<S: Scale, const BITS: u32, const BLOCK: usize>(
15    values: &[f32],
16) -> Result<Quantized<S>> {
17    quantize_with::<S>(values, BITS, BLOCK)
18}
19
20/// Runtime-width variant of [`quantize`].
21pub fn quantize_with<S: Scale>(values: &[f32], bits: u32, block: usize) -> Result<Quantized<S>> {
22    check_bits(bits)?;
23    check_block(block)?;
24    if values.is_empty() {
25        return Ok(Quantized::Symmetric {
26            scales: Vec::new(),
27            codes: Packed::from_raw(Vec::new(), bits, 0),
28            block,
29            len: 0,
30        });
31    }
32    let (scales_f, codes) = quantize_sym_packed(values, bits, block);
33    Ok(Quantized::Symmetric {
34        scales: scales_f.into_iter().map(S::from_f32).collect(),
35        codes,
36        block,
37        len: values.len(),
38    })
39}
40
41/// Quantize the entire tensor with a single scale.
42pub fn quantize_tensor<S: Scale, const BITS: u32>(values: &[f32]) -> Result<Quantized<S>> {
43    quantize_with::<S>(values, BITS, values.len().max(1))
44}
45
46#[cfg(test)]
47mod tests {
48    use super::*;
49
50    #[test]
51    fn eight_bit_roundtrip_stays_within_half_step() {
52        let w = [0.42_f32, -0.10, 0.70, -0.50];
53        let q = quantize::<f32, 8, 4>(&w).unwrap();
54        let back = q.dequantize();
55        for (a, b) in w.iter().zip(&back) {
56            assert!((a - b).abs() < 0.01, "{a} vs {b}");
57        }
58    }
59
60    #[test]
61    fn packed_four_bit_uses_half_byte_per_code() {
62        let w = [0.1_f32; 32];
63        let q = quantize::<f32, 4, 32>(&w).unwrap();
64        assert_eq!(q.codes().len(), 16);
65    }
66
67    #[test]
68    fn remainder_block_roundtrips() {
69        let w: Vec<f32> = (0..40).map(|i| (i as f32) * 0.01 - 0.2).collect();
70        let q = quantize::<f32, 8, 32>(&w).unwrap();
71        assert_eq!(q.len(), 40);
72        let back = q.dequantize();
73        for (a, b) in w.iter().zip(&back) {
74            assert!((a - b).abs() < 0.01, "{a} vs {b}");
75        }
76    }
77
78    #[test]
79    fn dequantize_into_rejects_wrong_length() {
80        let w = [0.1_f32; 8];
81        let q = quantize::<f32, 8, 8>(&w).unwrap();
82        let mut out = [0.0f32; 3];
83        assert!(matches!(
84            q.dequantize_into(&mut out),
85            Err(crate::Error::LengthMismatch {
86                expected: 8,
87                got: 3
88            })
89        ));
90    }
91
92    #[test]
93    fn four_bit_remainder_roundtrips() {
94        let w: Vec<f32> = (0..40).map(|i| (i as f32) * 0.02 - 0.4).collect();
95        let q = quantize::<f32, 4, 32>(&w).unwrap();
96        let back = q.dequantize();
97        for (a, b) in w.iter().zip(&back) {
98            assert!((a - b).abs() < 0.08, "{a} vs {b}");
99        }
100    }
101
102    #[test]
103    fn fused_dot_matches_dequant_then_dot() {
104        let w: Vec<f32> = (0..64).map(|i| (i as f32) * 0.01 - 0.3).collect();
105        let q = quantize::<f32, 8, 32>(&w).unwrap();
106        let recon = q.dequantize();
107        let naive: f32 = recon.iter().zip(&w).map(|(a, b)| a * b).sum();
108        let fused = q.dot(&w).unwrap();
109        assert!((naive - fused).abs() < 1e-4, "{naive} vs {fused}");
110    }
111}
quantize/methods/symmetric/mod.rs

quantize/methods/symmetric/
mod.rs
