ipfrs_tensorlogic/
quantizer.rs

1//! TensorQuantizer — symmetric and asymmetric quantization for tensor values.
2//!
3//! Provides calibration-based quantization with INT8 and INT4 support,
4//! including quantize/dequantize roundtrips and MSE error measurement.
5//!
6//! # Examples
7//!
8//! ```
9//! use ipfrs_tensorlogic::quantizer::{TensorQuantizer, QuantMode, QuantBits};
10//!
11//! let mut quantizer = TensorQuantizer::new();
12//! quantizer.calibrate(&[1.0, -0.5, 0.3, -1.0]);
13//!
14//! let params = quantizer.compute_params(QuantMode::Symmetric, QuantBits::Int8).expect("example: should succeed in docs");
15//! let quantized = TensorQuantizer::quantize(&[0.5, -0.5], &params);
16//! let restored = TensorQuantizer::dequantize(&quantized, &params);
17//! ```
18
19/// Quantization mode.
20#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21pub enum QuantMode {
22    /// Symmetric: zero_point = 0, scale = max(|min|, |max|) / (2^(bits-1) - 1)
23    Symmetric,
24    /// Asymmetric: scale = (max - min) / (2^bits - 1), zero_point = round(-min / scale)
25    Asymmetric,
26}
27
28/// Quantization bit width.
29#[derive(Debug, Clone, Copy, PartialEq, Eq)]
30pub enum QuantBits {
31    /// 8-bit: symmetric range [-128, 127], asymmetric range [0, 255]
32    Int8,
33    /// 4-bit: symmetric range [-8, 7], asymmetric range [0, 15]
34    Int4,
35}
36
37impl QuantBits {
38    /// Number of bits for this quantization level.
39    fn num_bits(self) -> u32 {
40        match self {
41            Self::Int8 => 8,
42            Self::Int4 => 4,
43        }
44    }
45
46    /// Signed symmetric range: [-(2^(bits-1)), 2^(bits-1) - 1]
47    fn symmetric_range(self) -> (i32, i32) {
48        let half = 1_i32 << (self.num_bits() - 1);
49        (-half, half - 1)
50    }
51
52    /// Unsigned asymmetric range: [0, 2^bits - 1]
53    fn asymmetric_range(self) -> (i32, i32) {
54        let max = (1_i32 << self.num_bits()) - 1;
55        (0, max)
56    }
57}
58
59/// Parameters computed from calibration data for quantization/dequantization.
60#[derive(Debug, Clone)]
61pub struct QuantParams {
62    /// Scale factor: maps floating-point range to integer range.
63    pub scale: f64,
64    /// Zero point offset (0 for symmetric mode).
65    pub zero_point: i32,
66    /// Quantization mode used.
67    pub mode: QuantMode,
68    /// Bit width used.
69    pub bits: QuantBits,
70}
71
72/// Statistics about the quantizer's calibration state.
73#[derive(Debug, Clone)]
74pub struct QuantizerStats {
75    /// Number of samples observed during calibration.
76    pub samples_seen: u64,
77    /// Minimum value observed.
78    pub calibration_min: f64,
79    /// Maximum value observed.
80    pub calibration_max: f64,
81}
82
83/// Calibration-based tensor quantizer.
84///
85/// Usage: call [`calibrate`](TensorQuantizer::calibrate) with representative
86/// data, then [`compute_params`](TensorQuantizer::compute_params) to derive
87/// quantization parameters.
88pub struct TensorQuantizer {
89    calibration_min: f64,
90    calibration_max: f64,
91    samples_seen: u64,
92}
93
94impl TensorQuantizer {
95    /// Create a new uncalibrated quantizer.
96    pub fn new() -> Self {
97        Self {
98            calibration_min: f64::MAX,
99            calibration_max: f64::MIN,
100            samples_seen: 0,
101        }
102    }
103
104    /// Update calibration range from observed values.
105    ///
106    /// Can be called multiple times; min/max accumulate across calls.
107    pub fn calibrate(&mut self, values: &[f64]) {
108        for &v in values {
109            if v < self.calibration_min {
110                self.calibration_min = v;
111            }
112            if v > self.calibration_max {
113                self.calibration_max = v;
114            }
115        }
116        self.samples_seen += values.len() as u64;
117    }
118
119    /// Compute quantization parameters from calibration data.
120    ///
121    /// Returns an error if no samples have been observed.
122    pub fn compute_params(&self, mode: QuantMode, bits: QuantBits) -> Result<QuantParams, String> {
123        if self.samples_seen == 0 {
124            return Err("No calibration data: call calibrate() first".to_string());
125        }
126
127        let (scale, zero_point) = match mode {
128            QuantMode::Symmetric => {
129                let abs_max = self.calibration_min.abs().max(self.calibration_max.abs());
130                let (_qmin, qmax) = bits.symmetric_range();
131                let s = if abs_max == 0.0 {
132                    1.0
133                } else {
134                    abs_max / qmax as f64
135                };
136                (s, 0)
137            }
138            QuantMode::Asymmetric => {
139                let range = self.calibration_max - self.calibration_min;
140                let (_qmin, qmax) = bits.asymmetric_range();
141                let s = if range == 0.0 {
142                    1.0
143                } else {
144                    range / qmax as f64
145                };
146                let zp = (-self.calibration_min / s).round() as i32;
147                (s, zp)
148            }
149        };
150
151        Ok(QuantParams {
152            scale,
153            zero_point,
154            mode,
155            bits,
156        })
157    }
158
159    /// Quantize floating-point values to integers using the given parameters.
160    ///
161    /// Each value is mapped as: clamp(round(value / scale) + zero_point, qmin, qmax)
162    pub fn quantize(values: &[f64], params: &QuantParams) -> Vec<i32> {
163        let (qmin, qmax) = match params.mode {
164            QuantMode::Symmetric => params.bits.symmetric_range(),
165            QuantMode::Asymmetric => params.bits.asymmetric_range(),
166        };
167
168        values
169            .iter()
170            .map(|&v| {
171                let q = (v / params.scale).round() as i32 + params.zero_point;
172                q.clamp(qmin, qmax)
173            })
174            .collect()
175    }
176
177    /// Dequantize integer values back to floating-point.
178    ///
179    /// Each value is mapped as: (q - zero_point) * scale
180    pub fn dequantize(quantized: &[i32], params: &QuantParams) -> Vec<f64> {
181        quantized
182            .iter()
183            .map(|&q| (q - params.zero_point) as f64 * params.scale)
184            .collect()
185    }
186
187    /// Compute mean squared error between original values and their
188    /// quantize→dequantize roundtrip.
189    pub fn quantization_error(original: &[f64], params: &QuantParams) -> f64 {
190        if original.is_empty() {
191            return 0.0;
192        }
193        let quantized = Self::quantize(original, params);
194        let dequantized = Self::dequantize(&quantized, params);
195        let sum_sq: f64 = original
196            .iter()
197            .zip(dequantized.iter())
198            .map(|(&o, &d)| {
199                let diff = o - d;
200                diff * diff
201            })
202            .sum();
203        sum_sq / original.len() as f64
204    }
205
206    /// Reset calibration state.
207    pub fn reset_calibration(&mut self) {
208        self.calibration_min = f64::MAX;
209        self.calibration_max = f64::MIN;
210        self.samples_seen = 0;
211    }
212
213    /// Return current calibration statistics.
214    pub fn stats(&self) -> QuantizerStats {
215        QuantizerStats {
216            samples_seen: self.samples_seen,
217            calibration_min: self.calibration_min,
218            calibration_max: self.calibration_max,
219        }
220    }
221}
222
223impl Default for TensorQuantizer {
224    fn default() -> Self {
225        Self::new()
226    }
227}
228
229#[cfg(test)]
230mod tests {
231    use super::*;
232
233    // ── Symmetric INT8 ──────────────────────────────────────────────
234
235    #[test]
236    fn symmetric_int8_basic() {
237        let mut q = TensorQuantizer::new();
238        q.calibrate(&[-1.0, 0.0, 1.0]);
239        let params = q
240            .compute_params(QuantMode::Symmetric, QuantBits::Int8)
241            .expect("params");
242        assert_eq!(params.zero_point, 0);
243        // scale = 1.0 / 127
244        let expected_scale = 1.0 / 127.0;
245        assert!((params.scale - expected_scale).abs() < 1e-12);
246    }
247
248    #[test]
249    fn symmetric_int8_quantize_dequantize() {
250        let mut q = TensorQuantizer::new();
251        q.calibrate(&[-1.0, 1.0]);
252        let params = q
253            .compute_params(QuantMode::Symmetric, QuantBits::Int8)
254            .expect("params");
255
256        let values = vec![0.0, 0.5, -0.5, 1.0, -1.0];
257        let quantized = TensorQuantizer::quantize(&values, &params);
258        let dequantized = TensorQuantizer::dequantize(&quantized, &params);
259
260        for (&orig, &deq) in values.iter().zip(dequantized.iter()) {
261            assert!((orig - deq).abs() < 0.01, "orig={orig} deq={deq}");
262        }
263    }
264
265    #[test]
266    fn symmetric_int8_zero_point_is_zero() {
267        let mut q = TensorQuantizer::new();
268        q.calibrate(&[-5.0, 3.0]);
269        let params = q
270            .compute_params(QuantMode::Symmetric, QuantBits::Int8)
271            .expect("params");
272        assert_eq!(params.zero_point, 0);
273    }
274
275    #[test]
276    fn symmetric_int8_large_range() {
277        let mut q = TensorQuantizer::new();
278        q.calibrate(&[-100.0, 100.0]);
279        let params = q
280            .compute_params(QuantMode::Symmetric, QuantBits::Int8)
281            .expect("params");
282        let expected_scale = 100.0 / 127.0;
283        assert!((params.scale - expected_scale).abs() < 1e-10);
284    }
285
286    // ── Asymmetric INT8 ─────────────────────────────────────────────
287
288    #[test]
289    fn asymmetric_int8_basic() {
290        let mut q = TensorQuantizer::new();
291        q.calibrate(&[0.0, 1.0]);
292        let params = q
293            .compute_params(QuantMode::Asymmetric, QuantBits::Int8)
294            .expect("params");
295        // scale = 1.0 / 255
296        let expected_scale = 1.0 / 255.0;
297        assert!((params.scale - expected_scale).abs() < 1e-12);
298        // zero_point = round(0 / scale) = 0
299        assert_eq!(params.zero_point, 0);
300    }
301
302    #[test]
303    fn asymmetric_int8_negative_range() {
304        let mut q = TensorQuantizer::new();
305        q.calibrate(&[-2.0, 2.0]);
306        let params = q
307            .compute_params(QuantMode::Asymmetric, QuantBits::Int8)
308            .expect("params");
309        // scale = 4.0 / 255
310        let expected_scale = 4.0 / 255.0;
311        assert!((params.scale - expected_scale).abs() < 1e-10);
312        // zero_point = round(2.0 / scale) = round(127.5) = 128
313        let expected_zp = (2.0 / expected_scale).round() as i32;
314        assert_eq!(params.zero_point, expected_zp);
315    }
316
317    #[test]
318    fn asymmetric_int8_roundtrip() {
319        let mut q = TensorQuantizer::new();
320        q.calibrate(&[-1.0, 3.0]);
321        let params = q
322            .compute_params(QuantMode::Asymmetric, QuantBits::Int8)
323            .expect("params");
324        let values = vec![0.0, 1.0, 2.0, 3.0, -1.0];
325        let quantized = TensorQuantizer::quantize(&values, &params);
326        let dequantized = TensorQuantizer::dequantize(&quantized, &params);
327        for (&orig, &deq) in values.iter().zip(dequantized.iter()) {
328            assert!((orig - deq).abs() < 0.05, "orig={orig} deq={deq}");
329        }
330    }
331
332    // ── INT4 quantization ───────────────────────────────────────────
333
334    #[test]
335    fn symmetric_int4_basic() {
336        let mut q = TensorQuantizer::new();
337        q.calibrate(&[-1.0, 1.0]);
338        let params = q
339            .compute_params(QuantMode::Symmetric, QuantBits::Int4)
340            .expect("params");
341        assert_eq!(params.zero_point, 0);
342        // scale = 1.0 / 7
343        let expected_scale = 1.0 / 7.0;
344        assert!((params.scale - expected_scale).abs() < 1e-12);
345    }
346
347    #[test]
348    fn symmetric_int4_clamping() {
349        let mut q = TensorQuantizer::new();
350        q.calibrate(&[-1.0, 1.0]);
351        let params = q
352            .compute_params(QuantMode::Symmetric, QuantBits::Int4)
353            .expect("params");
354        // Value well outside range should clamp
355        let quantized = TensorQuantizer::quantize(&[10.0], &params);
356        assert_eq!(quantized[0], 7); // clamped to qmax
357        let quantized_neg = TensorQuantizer::quantize(&[-10.0], &params);
358        assert_eq!(quantized_neg[0], -8); // clamped to qmin
359    }
360
361    #[test]
362    fn asymmetric_int4_basic() {
363        let mut q = TensorQuantizer::new();
364        q.calibrate(&[0.0, 1.0]);
365        let params = q
366            .compute_params(QuantMode::Asymmetric, QuantBits::Int4)
367            .expect("params");
368        // scale = 1.0 / 15
369        let expected_scale = 1.0 / 15.0;
370        assert!((params.scale - expected_scale).abs() < 1e-12);
371        assert_eq!(params.zero_point, 0);
372    }
373
374    #[test]
375    fn asymmetric_int4_roundtrip() {
376        let mut q = TensorQuantizer::new();
377        q.calibrate(&[-2.0, 2.0]);
378        let params = q
379            .compute_params(QuantMode::Asymmetric, QuantBits::Int4)
380            .expect("params");
381        let values = vec![-2.0, -1.0, 0.0, 1.0, 2.0];
382        let quantized = TensorQuantizer::quantize(&values, &params);
383        let dequantized = TensorQuantizer::dequantize(&quantized, &params);
384        for (&orig, &deq) in values.iter().zip(dequantized.iter()) {
385            // INT4 has coarser granularity
386            assert!((orig - deq).abs() < 0.5, "orig={orig} deq={deq}");
387        }
388    }
389
390    // ── Calibration accumulation ────────────────────────────────────
391
392    #[test]
393    fn calibration_accumulates_across_calls() {
394        let mut q = TensorQuantizer::new();
395        q.calibrate(&[0.0, 1.0]);
396        q.calibrate(&[-2.0, 0.5]);
397        q.calibrate(&[0.0, 3.0]);
398        let stats = q.stats();
399        assert_eq!(stats.samples_seen, 6);
400        assert!((stats.calibration_min - (-2.0)).abs() < 1e-15);
401        assert!((stats.calibration_max - 3.0).abs() < 1e-15);
402    }
403
404    #[test]
405    fn calibration_single_value() {
406        let mut q = TensorQuantizer::new();
407        q.calibrate(&[5.0]);
408        let stats = q.stats();
409        assert_eq!(stats.samples_seen, 1);
410        assert!((stats.calibration_min - 5.0).abs() < 1e-15);
411        assert!((stats.calibration_max - 5.0).abs() < 1e-15);
412    }
413
414    // ── Roundtrip error ─────────────────────────────────────────────
415
416    #[test]
417    fn roundtrip_error_is_small_int8() {
418        let mut q = TensorQuantizer::new();
419        let values: Vec<f64> = (0..100).map(|i| (i as f64 - 50.0) / 50.0).collect();
420        q.calibrate(&values);
421        let params = q
422            .compute_params(QuantMode::Symmetric, QuantBits::Int8)
423            .expect("params");
424        let mse = TensorQuantizer::quantization_error(&values, &params);
425        assert!(mse < 0.001, "MSE too large: {mse}");
426    }
427
428    #[test]
429    fn roundtrip_error_larger_for_int4() {
430        let mut q = TensorQuantizer::new();
431        let values: Vec<f64> = (0..100).map(|i| (i as f64 - 50.0) / 50.0).collect();
432        q.calibrate(&values);
433        let params_8 = q
434            .compute_params(QuantMode::Symmetric, QuantBits::Int8)
435            .expect("params8");
436        let params_4 = q
437            .compute_params(QuantMode::Symmetric, QuantBits::Int4)
438            .expect("params4");
439        let mse_8 = TensorQuantizer::quantization_error(&values, &params_8);
440        let mse_4 = TensorQuantizer::quantization_error(&values, &params_4);
441        assert!(mse_4 > mse_8, "INT4 error should exceed INT8 error");
442    }
443
444    // ── quantization_error MSE correctness ──────────────────────────
445
446    #[test]
447    fn quantization_error_manual_check() {
448        let mut q = TensorQuantizer::new();
449        q.calibrate(&[-1.0, 1.0]);
450        let params = q
451            .compute_params(QuantMode::Symmetric, QuantBits::Int8)
452            .expect("params");
453        let original = vec![0.5];
454        let quantized = TensorQuantizer::quantize(&original, &params);
455        let dequantized = TensorQuantizer::dequantize(&quantized, &params);
456        let diff = original[0] - dequantized[0];
457        let expected_mse = diff * diff;
458        let mse = TensorQuantizer::quantization_error(&original, &params);
459        assert!((mse - expected_mse).abs() < 1e-15);
460    }
461
462    #[test]
463    fn quantization_error_empty_input() {
464        let mut q = TensorQuantizer::new();
465        q.calibrate(&[1.0]);
466        let params = q
467            .compute_params(QuantMode::Symmetric, QuantBits::Int8)
468            .expect("params");
469        let mse = TensorQuantizer::quantization_error(&[], &params);
470        assert!((mse - 0.0).abs() < 1e-15);
471    }
472
473    // ── Edge cases ──────────────────────────────────────────────────
474
475    #[test]
476    fn all_zeros() {
477        let mut q = TensorQuantizer::new();
478        q.calibrate(&[0.0, 0.0, 0.0]);
479        let params = q
480            .compute_params(QuantMode::Symmetric, QuantBits::Int8)
481            .expect("params");
482        // scale should be 1.0 (fallback for zero range)
483        assert!((params.scale - 1.0).abs() < 1e-15);
484        let quantized = TensorQuantizer::quantize(&[0.0, 0.0], &params);
485        assert!(quantized.iter().all(|&q| q == 0));
486    }
487
488    #[test]
489    fn all_zeros_asymmetric() {
490        let mut q = TensorQuantizer::new();
491        q.calibrate(&[0.0, 0.0]);
492        let params = q
493            .compute_params(QuantMode::Asymmetric, QuantBits::Int8)
494            .expect("params");
495        assert!((params.scale - 1.0).abs() < 1e-15);
496    }
497
498    #[test]
499    fn single_value_symmetric() {
500        let mut q = TensorQuantizer::new();
501        q.calibrate(&[5.0]);
502        let params = q
503            .compute_params(QuantMode::Symmetric, QuantBits::Int8)
504            .expect("params");
505        let expected_scale = 5.0 / 127.0;
506        assert!((params.scale - expected_scale).abs() < 1e-12);
507    }
508
509    #[test]
510    fn negative_only_values() {
511        let mut q = TensorQuantizer::new();
512        q.calibrate(&[-3.0, -1.0, -2.0]);
513        let params = q
514            .compute_params(QuantMode::Symmetric, QuantBits::Int8)
515            .expect("params");
516        assert_eq!(params.zero_point, 0);
517        let expected_scale = 3.0 / 127.0;
518        assert!((params.scale - expected_scale).abs() < 1e-12);
519    }
520
521    #[test]
522    fn negative_only_asymmetric() {
523        let mut q = TensorQuantizer::new();
524        q.calibrate(&[-3.0, -1.0]);
525        let params = q
526            .compute_params(QuantMode::Asymmetric, QuantBits::Int8)
527            .expect("params");
528        // scale = 2.0 / 255
529        let expected_scale = 2.0 / 255.0;
530        assert!((params.scale - expected_scale).abs() < 1e-10);
531        // zero_point = round(3.0 / scale)
532        let expected_zp = (3.0 / expected_scale).round() as i32;
533        assert_eq!(params.zero_point, expected_zp);
534    }
535
536    // ── Reset calibration ───────────────────────────────────────────
537
538    #[test]
539    fn reset_calibration_clears_state() {
540        let mut q = TensorQuantizer::new();
541        q.calibrate(&[1.0, 2.0, 3.0]);
542        q.reset_calibration();
543        let stats = q.stats();
544        assert_eq!(stats.samples_seen, 0);
545        assert_eq!(stats.calibration_min, f64::MAX);
546        assert_eq!(stats.calibration_max, f64::MIN);
547    }
548
549    #[test]
550    fn reset_then_recalibrate() {
551        let mut q = TensorQuantizer::new();
552        q.calibrate(&[-10.0, 10.0]);
553        q.reset_calibration();
554        q.calibrate(&[-1.0, 1.0]);
555        let params = q
556            .compute_params(QuantMode::Symmetric, QuantBits::Int8)
557            .expect("params");
558        let expected_scale = 1.0 / 127.0;
559        assert!((params.scale - expected_scale).abs() < 1e-12);
560    }
561
562    // ── Uncalibrated error ──────────────────────────────────────────
563
564    #[test]
565    fn error_on_uncalibrated() {
566        let q = TensorQuantizer::new();
567        let result = q.compute_params(QuantMode::Symmetric, QuantBits::Int8);
568        assert!(result.is_err());
569        assert!(result.expect_err("should be err").contains("calibration"));
570    }
571
572    // ── Clamping at boundaries ──────────────────────────────────────
573
574    #[test]
575    fn clamping_symmetric_int8() {
576        let mut q = TensorQuantizer::new();
577        q.calibrate(&[-1.0, 1.0]);
578        let params = q
579            .compute_params(QuantMode::Symmetric, QuantBits::Int8)
580            .expect("params");
581        // Values far beyond calibration range should clamp
582        let quantized = TensorQuantizer::quantize(&[1000.0, -1000.0], &params);
583        assert_eq!(quantized[0], 127);
584        assert_eq!(quantized[1], -128);
585    }
586
587    #[test]
588    fn clamping_asymmetric_int8() {
589        let mut q = TensorQuantizer::new();
590        q.calibrate(&[0.0, 1.0]);
591        let params = q
592            .compute_params(QuantMode::Asymmetric, QuantBits::Int8)
593            .expect("params");
594        let quantized = TensorQuantizer::quantize(&[1000.0, -1000.0], &params);
595        assert_eq!(quantized[0], 255);
596        assert_eq!(quantized[1], 0);
597    }
598
599    // ── Default trait ───────────────────────────────────────────────
600
601    #[test]
602    fn default_trait_works() {
603        let q = TensorQuantizer::default();
604        let stats = q.stats();
605        assert_eq!(stats.samples_seen, 0);
606    }
607
608    // ── QuantBits helpers ───────────────────────────────────────────
609
610    #[test]
611    fn quant_bits_ranges() {
612        assert_eq!(QuantBits::Int8.symmetric_range(), (-128, 127));
613        assert_eq!(QuantBits::Int8.asymmetric_range(), (0, 255));
614        assert_eq!(QuantBits::Int4.symmetric_range(), (-8, 7));
615        assert_eq!(QuantBits::Int4.asymmetric_range(), (0, 15));
616    }
617
618    // ── Stats ───────────────────────────────────────────────────────
619
620    #[test]
621    fn stats_reflect_calibration() {
622        let mut q = TensorQuantizer::new();
623        q.calibrate(&[1.0, 2.0]);
624        q.calibrate(&[3.0]);
625        let stats = q.stats();
626        assert_eq!(stats.samples_seen, 3);
627        assert!((stats.calibration_min - 1.0).abs() < 1e-15);
628        assert!((stats.calibration_max - 3.0).abs() < 1e-15);
629    }
630
631    // ── Mixed mode comparison ───────────────────────────────────────
632
633    #[test]
634    fn symmetric_vs_asymmetric_error() {
635        let mut q = TensorQuantizer::new();
636        // Asymmetric data: all positive — asymmetric should use range better
637        let values: Vec<f64> = (0..50).map(|i| i as f64 / 50.0).collect();
638        q.calibrate(&values);
639        let sym_params = q
640            .compute_params(QuantMode::Symmetric, QuantBits::Int8)
641            .expect("sym");
642        let asym_params = q
643            .compute_params(QuantMode::Asymmetric, QuantBits::Int8)
644            .expect("asym");
645        let sym_err = TensorQuantizer::quantization_error(&values, &sym_params);
646        let asym_err = TensorQuantizer::quantization_error(&values, &asym_params);
647        // For all-positive data, asymmetric should have equal or less error
648        assert!(
649            asym_err <= sym_err + 1e-10,
650            "asym_err={asym_err} sym_err={sym_err}"
651        );
652    }
653}
ipfrs_tensorlogic/quantizer.rs

ipfrs_tensorlogic/
quantizer.rs
