teeny_quant/quant/
affine.rs

1/*
2 * Copyright (c) 2026 Teenygrad.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *   http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17//! Affine (scale + zero-point) integer quantization, shared by the INT8 and INT4 schemes --
18//! INT4 just runs this with `bits: 4` and then nibble-packs the result (see
19//! [`crate::quant::pack4`]).
20
21use crate::error::{Error, Result};
22use crate::quant::granularity::Granularity;
23use crate::quant::groups::assign_groups;
24
25/// Per-group affine quantization parameters: `q = round(x / scale) + zero_point`, clamped to the
26/// scheme's `[qmin, qmax]`.
27#[derive(Debug, Clone, Copy, PartialEq)]
28pub struct AffineParams {
29    /// The quantization scale (always > 0).
30    pub scale: f32,
31    /// The zero-point in quantized (integer) units. `0` for symmetric quantization.
32    pub zero_point: i32,
33}
34
35/// A tensor quantized with one [`AffineParams`] per group.
36#[derive(Debug, Clone)]
37pub struct QuantizedAffine {
38    /// One quantized value per input element, in the same (row-major) order as the input.
39    pub qvalues: Vec<i8>,
40    /// One [`AffineParams`] per group.
41    pub params: Vec<AffineParams>,
42    /// `qvalues[i]`'s group is `groups[i]`, i.e. `params[groups[i]]` quantized it.
43    pub groups: Vec<u32>,
44}
45
46/// The `[qmin, qmax]` integer range for `bits`-bit (signed) quantization.
47fn qrange(bits: u8, symmetric: bool) -> (i32, i32) {
48    let qmax = (1i32 << (bits - 1)) - 1; // 127 for bits=8, 7 for bits=4
49    let qmin = if symmetric { -qmax } else { -qmax - 1 }; // -127/-128, -7/-8
50    (qmin, qmax)
51}
52
53/// Quantizes `data` (row-major, shape `shape`) to `bits`-bit signed integers at the given
54/// `granularity`, computing one [`AffineParams`] per group from that group's own min/max.
55pub fn quantize_affine(
56    tensor: &str,
57    data: &[f32],
58    shape: &[usize],
59    granularity: Granularity,
60    symmetric: bool,
61    bits: u8,
62) -> Result<QuantizedAffine> {
63    let n: usize = shape.iter().product();
64    if data.len() != n {
65        return Err(Error::ShapeMismatch {
66            name: tensor.to_string(),
67            a: vec![data.len()],
68            b: shape.to_vec(),
69        });
70    }
71
72    let (qmin, qmax) = qrange(bits, symmetric);
73    let axis_group = granularity.axis_and_group_size();
74
75    if let Some((axis, group_size)) = axis_group {
76        if axis >= shape.len() {
77            return Err(Error::InvalidAxis {
78                tensor: tensor.to_string(),
79                axis,
80                rank: shape.len(),
81            });
82        }
83        if group_size == 0 {
84            return Err(Error::InvalidGroupSize(tensor.to_string()));
85        }
86    }
87
88    let (groups, ngroups) = assign_groups(shape, granularity);
89
90    // Pass 1: per-group min/max (and amax, for symmetric).
91    let mut vmin = vec![f32::INFINITY; ngroups];
92    let mut vmax = vec![f32::NEG_INFINITY; ngroups];
93    for (i, &x) in data.iter().enumerate() {
94        let g = groups[i] as usize;
95        if x < vmin[g] {
96            vmin[g] = x;
97        }
98        if x > vmax[g] {
99            vmax[g] = x;
100        }
101    }
102
103    let params: Vec<AffineParams> = (0..ngroups)
104        .map(|g| {
105            let (lo, hi) = if vmin[g].is_finite() {
106                (vmin[g].min(0.0), vmax[g].max(0.0))
107            } else {
108                (0.0, 0.0) // empty group (shouldn't happen, but stay total)
109            };
110            if symmetric {
111                let amax = lo.abs().max(hi.abs());
112                let scale = if amax == 0.0 { 1.0 } else { amax / qmax as f32 };
113                AffineParams {
114                    scale,
115                    zero_point: 0,
116                }
117            } else {
118                let scale = if hi == lo {
119                    1.0
120                } else {
121                    (hi - lo) / (qmax - qmin) as f32
122                };
123                let zero_point = (qmin as f32 - lo / scale).round() as i32;
124                AffineParams {
125                    scale,
126                    zero_point: zero_point.clamp(qmin, qmax),
127                }
128            }
129        })
130        .collect();
131
132    // Pass 2: quantize each element with its group's params.
133    let qvalues: Vec<i8> = data
134        .iter()
135        .enumerate()
136        .map(|(i, &x)| {
137            let p = params[groups[i] as usize];
138            let q = (x / p.scale).round() as i32 + p.zero_point;
139            q.clamp(qmin, qmax) as i8
140        })
141        .collect();
142
143    Ok(QuantizedAffine {
144        qvalues,
145        params,
146        groups,
147    })
148}
149
150/// Reconstructs `f32` values from `qvalues`/`params`/`groups` produced by [`quantize_affine`].
151pub fn dequantize_affine(q: &QuantizedAffine) -> Vec<f32> {
152    q.qvalues
153        .iter()
154        .enumerate()
155        .map(|(i, &v)| {
156            let p = q.params[q.groups[i] as usize];
157            (v as i32 - p.zero_point) as f32 * p.scale
158        })
159        .collect()
160}
161
162#[cfg(test)]
163mod tests {
164    use super::*;
165
166    #[test]
167    fn per_tensor_symmetric_round_trip() {
168        let data = vec![-1.0f32, -0.5, 0.0, 0.5, 1.0, 2.0];
169        let shape = [6usize];
170        let q = quantize_affine("t", &data, &shape, Granularity::PerTensor, true, 8).unwrap();
171        assert_eq!(q.params.len(), 1);
172        assert_eq!(q.params[0].zero_point, 0);
173        let deq = dequantize_affine(&q);
174        for (orig, back) in data.iter().zip(deq.iter()) {
175            assert!((orig - back).abs() < 0.02, "{orig} vs {back}");
176        }
177    }
178
179    #[test]
180    fn per_tensor_asymmetric_covers_positive_only_range() {
181        // All-positive data: asymmetric should use the full qmin..qmax range, unlike symmetric.
182        let data = vec![0.0f32, 1.0, 2.0, 3.0, 4.0];
183        let shape = [5usize];
184        let q = quantize_affine("t", &data, &shape, Granularity::PerTensor, false, 8).unwrap();
185        let deq = dequantize_affine(&q);
186        for (orig, back) in data.iter().zip(deq.iter()) {
187            assert!((orig - back).abs() < 0.05, "{orig} vs {back}");
188        }
189    }
190
191    #[test]
192    fn per_channel_independent_scales() {
193        // shape [2, 4]: row 0 has small magnitude, row 1 much larger -- per-channel (axis 0)
194        // should give each row its own scale, so both round-trip tightly.
195        let data = vec![0.1f32, -0.1, 0.05, -0.05, 100.0, -100.0, 50.0, -50.0];
196        let shape = [2usize, 4];
197        let q = quantize_affine(
198            "t",
199            &data,
200            &shape,
201            Granularity::PerChannel { axis: 0 },
202            true,
203            8,
204        )
205        .unwrap();
206        assert_eq!(q.params.len(), 2);
207        let deq = dequantize_affine(&q);
208        for (orig, back) in data.iter().zip(deq.iter()) {
209            let tol = orig.abs().max(1.0) * 0.02;
210            assert!((orig - back).abs() < tol, "{orig} vs {back}");
211        }
212    }
213
214    #[test]
215    fn group_wise_matches_per_channel_when_group_size_covers_axis() {
216        let data: Vec<f32> = (0..8).map(|i| i as f32 - 4.0).collect();
217        let shape = [2usize, 4];
218        let per_channel = quantize_affine(
219            "t",
220            &data,
221            &shape,
222            Granularity::PerChannel { axis: 0 },
223            true,
224            8,
225        )
226        .unwrap();
227        let group = quantize_affine(
228            "t",
229            &data,
230            &shape,
231            Granularity::Group {
232                axis: 1,
233                group_size: 4,
234            },
235            true,
236            8,
237        )
238        .unwrap();
239        // Different axes but both reduce to "one group per row of 4" here, so params should
240        // come out identical (grouping along the full-width axis == per-channel over rows).
241        assert_eq!(per_channel.params.len(), group.params.len());
242    }
243
244    #[test]
245    fn rejects_zero_group_size() {
246        let data = vec![1.0f32; 4];
247        let shape = [4usize];
248        let err = quantize_affine(
249            "t",
250            &data,
251            &shape,
252            Granularity::Group {
253                axis: 0,
254                group_size: 0,
255            },
256            true,
257            8,
258        )
259        .unwrap_err();
260        assert!(matches!(err, Error::InvalidGroupSize(_)));
261    }
262
263    #[test]
264    fn rejects_out_of_range_axis() {
265        let data = vec![1.0f32; 4];
266        let shape = [4usize];
267        let err = quantize_affine(
268            "t",
269            &data,
270            &shape,
271            Granularity::PerChannel { axis: 3 },
272            true,
273            8,
274        )
275        .unwrap_err();
276        assert!(matches!(err, Error::InvalidAxis { .. }));
277    }
278
279    #[test]
280    fn int4_range_is_clamped() {
281        let data = vec![-100.0f32, 100.0];
282        let shape = [2usize];
283        let q = quantize_affine("t", &data, &shape, Granularity::PerTensor, true, 4).unwrap();
284        for &v in &q.qvalues {
285            assert!((-7..=7).contains(&v));
286        }
287    }
288}
teeny_quant/quant/affine.rs

teeny_quant/quant/
affine.rs
