ailake_vec/
quantize.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2use half::f16;
3
4#[derive(Debug, Clone, Copy)]
5pub struct ScalingParams {
6    pub scale: f32,
7    pub zero_point: f32,
8}
9
10pub struct Quantizer;
11
12impl Quantizer {
13    pub fn f32_to_f16_bytes(v: &[f32]) -> Vec<u8> {
14        let mut out = Vec::with_capacity(v.len() * 2);
15        for &x in v {
16            out.extend_from_slice(&f16::from_f32(x).to_le_bytes());
17        }
18        out
19    }
20
21    pub fn f16_bytes_to_f32(bytes: &[u8]) -> Vec<f32> {
22        bytes
23            .chunks_exact(2)
24            .map(|b| f16::from_le_bytes([b[0], b[1]]).to_f32())
25            .collect()
26    }
27
28    /// Reinterpret raw little-endian F32 bytes (4 bytes/element) — the identity
29    /// decode for `VectorPrecision::F32` columns, no quantization involved.
30    pub fn f32_bytes_to_f32(bytes: &[u8]) -> Vec<f32> {
31        bytes
32            .chunks_exact(4)
33            .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
34            .collect()
35    }
36
37    pub fn f32_to_i8(v: &[f32]) -> (Vec<i8>, ScalingParams) {
38        let min = v.iter().cloned().fold(f32::INFINITY, f32::min);
39        let max = v.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
40        let range = max - min;
41        let scale = if range == 0.0 { 1.0 } else { range / 254.0 };
42        let zero_point = -128.0 - min / scale;
43        let quant = v
44            .iter()
45            .map(|&x| ((x / scale + zero_point).round().clamp(-128.0, 127.0)) as i8)
46            .collect();
47        (quant, ScalingParams { scale, zero_point })
48    }
49
50    pub fn i8_to_f32(v: &[i8], params: &ScalingParams) -> Vec<f32> {
51        v.iter()
52            .map(|&x| (x as f32 - params.zero_point) * params.scale)
53            .collect()
54    }
55}
56
57#[cfg(test)]
58mod tests {
59    use super::*;
60
61    #[test]
62    fn f16_roundtrip() {
63        let original: Vec<f32> = vec![0.1, -0.5, 1.0, 0.0, 100.0];
64        let bytes = Quantizer::f32_to_f16_bytes(&original);
65        let decoded = Quantizer::f16_bytes_to_f32(&bytes);
66        for (a, b) in original.iter().zip(decoded.iter()) {
67            assert!((a - b).abs() < 0.01, "f16 roundtrip error: {a} vs {b}");
68        }
69    }
70
71    #[test]
72    fn i8_roundtrip() {
73        let original: Vec<f32> = vec![0.0, 0.25, 0.5, 0.75, 1.0];
74        let (quant, params) = Quantizer::f32_to_i8(&original);
75        let decoded = Quantizer::i8_to_f32(&quant, &params);
76        for (a, b) in original.iter().zip(decoded.iter()) {
77            assert!((a - b).abs() < 0.02, "i8 roundtrip error: {a} vs {b}");
78        }
79    }
80}
ailake_vec/quantize.rs

ailake_vec/
quantize.rs
