quantize_rs/calibration/
methods.rs

1//! Calibration methods for quantization range optimization.
2
3use std::fmt;
4use std::str::FromStr;
5
6/// Strategy for choosing the quantization range from observed activations.
7///
8/// Marked `#[non_exhaustive]` so future calibration methods can be added
9/// without a breaking change.
10#[derive(Debug, Clone, Copy, Default)]
11#[non_exhaustive]
12pub enum CalibrationMethod {
13    /// Use the full observed min/max range (default).
14    #[default]
15    MinMax,
16
17    /// Clip outliers at the given percentile (e.g. 99.9).
18    Percentile(f32),
19
20    /// Minimize KL divergence between the original and quantized distributions.
21    Entropy,
22
23    /// Minimize mean squared error between original and dequantized values.
24    MSE,
25}
26
27impl fmt::Display for CalibrationMethod {
28    /// Renders the method in a form that round-trips through
29    /// [`FromStr`].  `Percentile(p)` is rendered as `"percentile:{p}"` so
30    /// the parameter survives the trip; the other variants render as their
31    /// lowercase keyword.
32    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
33        match self {
34            CalibrationMethod::MinMax => write!(f, "minmax"),
35            CalibrationMethod::Percentile(p) => write!(f, "percentile:{}", p),
36            CalibrationMethod::Entropy => write!(f, "entropy"),
37            CalibrationMethod::MSE => write!(f, "mse"),
38        }
39    }
40}
41
42impl FromStr for CalibrationMethod {
43    type Err = crate::errors::QuantizeError;
44
45    /// Parse a calibration method name.  Accepts:
46    ///   - `"minmax"`, `"entropy"`, `"mse"` (case-insensitive)
47    ///   - `"percentile"` (defaults to the 99.9th percentile)
48    ///   - `"percentile:NN"` for an explicit percentile, e.g. `"percentile:95"`
49    fn from_str(s: &str) -> Result<Self, Self::Err> {
50        let lower = s.to_lowercase();
51        if let Some(rest) = lower.strip_prefix("percentile:") {
52            let p: f32 = rest
53                .parse()
54                .map_err(|_| crate::errors::QuantizeError::Config {
55                    reason: format!(
56                        "Invalid percentile '{}'; expected a number like 'percentile:99.9'",
57                        rest
58                    ),
59                })?;
60            if !(0.0..=100.0).contains(&p) {
61                return Err(crate::errors::QuantizeError::Config {
62                    reason: format!("Percentile must be in [0, 100], got {}", p),
63                });
64            }
65            return Ok(CalibrationMethod::Percentile(p));
66        }
67        match lower.as_str() {
68            "minmax" => Ok(CalibrationMethod::MinMax),
69            "percentile" => Ok(CalibrationMethod::Percentile(99.9)),
70            "entropy" => Ok(CalibrationMethod::Entropy),
71            "mse" => Ok(CalibrationMethod::MSE),
72            _ => Err(crate::errors::QuantizeError::Config {
73                reason: format!(
74                    "Unknown calibration method: '{}'. Valid: minmax, percentile, percentile:N, entropy, mse",
75                    s
76                ),
77            }),
78        }
79    }
80}
81
82#[cfg(test)]
83mod tests {
84    use super::*;
85
86    #[test]
87    fn test_percentile_roundtrip() {
88        // Display → FromStr → Display preserves the percentile value.
89        let p1 = CalibrationMethod::Percentile(95.5);
90        let s = format!("{}", p1);
91        assert_eq!(s, "percentile:95.5");
92        let p2: CalibrationMethod = s.parse().unwrap();
93        assert!(matches!(p2, CalibrationMethod::Percentile(p) if (p - 95.5).abs() < 1e-6));
94    }
95
96    #[test]
97    fn test_bare_percentile_default() {
98        // "percentile" with no value falls back to 99.9 for back-compat.
99        let m: CalibrationMethod = "percentile".parse().unwrap();
100        assert!(matches!(m, CalibrationMethod::Percentile(p) if (p - 99.9).abs() < 1e-6));
101    }
102
103    #[test]
104    fn test_invalid_percentile_rejected() {
105        assert!("percentile:NaN_oops".parse::<CalibrationMethod>().is_err());
106        assert!("percentile:-1".parse::<CalibrationMethod>().is_err());
107        assert!("percentile:101".parse::<CalibrationMethod>().is_err());
108    }
109
110    #[test]
111    fn test_all_keywords_roundtrip() {
112        for method in [
113            CalibrationMethod::MinMax,
114            CalibrationMethod::Entropy,
115            CalibrationMethod::MSE,
116        ] {
117            let s = format!("{}", method);
118            let parsed: CalibrationMethod = s.parse().unwrap();
119            assert_eq!(
120                std::mem::discriminant(&method),
121                std::mem::discriminant(&parsed)
122            );
123        }
124    }
125}
quantize_rs/calibration/methods.rs

quantize_rs/calibration/
methods.rs
