1use crate::calibration::methods::CalibrationMethod;
7
8const NUM_BINS: usize = 256;
9
10#[derive(Debug, Clone)]
15pub struct ActivationStats {
16 min: f32,
17 max: f32,
18 mean: f32,
19 std: f32,
20 count: usize,
21
22 m2: f64,
24
25 histogram_bins: Vec<usize>,
26 hist_min: f32,
27 hist_max: f32,
28}
29
30impl ActivationStats {
31 pub fn min(&self) -> f32 {
33 self.min
34 }
35 pub fn max(&self) -> f32 {
37 self.max
38 }
39 pub fn mean(&self) -> f32 {
41 self.mean
42 }
43 pub fn std(&self) -> f32 {
45 self.std
46 }
47 pub fn count(&self) -> usize {
49 self.count
50 }
51}
52
53impl ActivationStats {
54 pub fn from_data(data: &[f32]) -> Self {
56 if data.is_empty() {
57 return Self::default();
58 }
59
60 let finite: Vec<f32> = data.iter().copied().filter(|v| v.is_finite()).collect();
61 if finite.is_empty() {
62 return Self::default();
63 }
64
65 let min = finite.iter().copied().fold(f32::INFINITY, f32::min);
66 let max = finite.iter().copied().fold(f32::NEG_INFINITY, f32::max);
67
68 let sum: f32 = finite.iter().sum();
69 let mean = sum / finite.len() as f32;
70
71 let m2: f64 = finite.iter().map(|&x| ((x - mean) as f64).powi(2)).sum();
72 let std = (m2 / finite.len() as f64).sqrt() as f32;
73
74 let histogram_bins = build_histogram(data, min, max);
75
76 Self {
77 min,
78 max,
79 mean,
80 std,
81 count: finite.len(),
82 m2,
83 histogram_bins,
84 hist_min: min,
85 hist_max: max,
86 }
87 }
88
89 pub fn update(&mut self, data: &[f32]) {
91 if data.is_empty() {
92 return;
93 }
94
95 let finite: Vec<f32> = data.iter().copied().filter(|v| v.is_finite()).collect();
97 if finite.is_empty() {
98 return;
99 }
100
101 if self.count == 0 {
111 *self = Self::from_data(data);
112 return;
113 }
114
115 let data_min = finite.iter().copied().fold(f32::INFINITY, f32::min);
116 let data_max = finite.iter().copied().fold(f32::NEG_INFINITY, f32::max);
117
118 let new_min = self.min.min(data_min);
119 let new_max = self.max.max(data_max);
120
121 let old_count = self.count as f64;
124 let new_count = finite.len() as f64;
125 let combined_count = old_count + new_count;
126
127 let data_sum: f64 = finite.iter().map(|&x| x as f64).sum();
128 let data_mean = data_sum / new_count;
129
130 let data_m2: f64 = finite
131 .iter()
132 .map(|&x| ((x as f64) - data_mean).powi(2))
133 .sum();
134
135 let delta = data_mean - self.mean as f64;
137 self.m2 = self.m2 + data_m2 + delta * delta * old_count * new_count / combined_count;
138
139 self.mean = (((self.mean as f64) * old_count + data_sum) / combined_count) as f32;
142 self.count = combined_count as usize;
143 self.std = (self.m2 / combined_count).sqrt() as f32;
144
145 if new_min < self.hist_min || new_max > self.hist_max {
147 let mut rebinned = vec![0usize; NUM_BINS];
148 rebin(
149 &self.histogram_bins,
150 self.hist_min,
151 self.hist_max,
152 &mut rebinned,
153 new_min,
154 new_max,
155 );
156 self.histogram_bins = rebinned;
157 self.hist_min = new_min;
158 self.hist_max = new_max;
159 }
160
161 let new_hist = build_histogram(&finite, self.hist_min, self.hist_max);
163 for (i, &c) in new_hist.iter().enumerate() {
164 self.histogram_bins[i] += c;
165 }
166
167 self.min = new_min;
168 self.max = new_max;
169 }
170
171 pub fn percentile(&self, p: f32) -> f32 {
173 if self.histogram_bins.is_empty() {
174 return self.min;
175 }
176
177 let total: usize = self.histogram_bins.iter().sum();
178 if total == 0 {
179 return self.min;
180 }
181
182 let target_count = (total as f32 * p / 100.0).ceil() as usize;
185 let mut cumulative = 0;
186
187 let bin_size = if (self.hist_max - self.hist_min).abs() < 1e-8 {
188 0.0
189 } else {
190 (self.hist_max - self.hist_min) / NUM_BINS as f32
191 };
192
193 for (i, &count) in self.histogram_bins.iter().enumerate() {
194 cumulative += count;
195 if cumulative >= target_count {
196 return self.hist_min + (i as f32 + 0.5) * bin_size;
197 }
198 }
199
200 self.max
201 }
202
203 pub fn histogram_data(&self) -> Vec<(f32, usize)> {
205 if (self.hist_max - self.hist_min).abs() < 1e-8 {
206 let total: usize = self.histogram_bins.iter().sum();
207 if total > 0 {
208 return vec![(self.hist_min, total)];
209 }
210 return Vec::new();
211 }
212 let bin_size = (self.hist_max - self.hist_min) / NUM_BINS as f32;
213 self.histogram_bins
214 .iter()
215 .enumerate()
216 .filter(|(_, &count)| count > 0)
217 .map(|(i, &count)| {
218 let value = self.hist_min + (i as f32 + 0.5) * bin_size;
219 (value, count)
220 })
221 .collect()
222 }
223}
224
225impl Default for ActivationStats {
226 fn default() -> Self {
227 Self {
228 min: f32::INFINITY,
229 max: f32::NEG_INFINITY,
230 mean: 0.0,
231 std: 0.0,
232 count: 0,
233 m2: 0.0,
234 histogram_bins: Vec::new(),
235 hist_min: 0.0,
236 hist_max: 0.0,
237 }
238 }
239}
240
241fn build_histogram(data: &[f32], min: f32, max: f32) -> Vec<usize> {
242 let mut bins = vec![0usize; NUM_BINS];
243
244 if (max - min).abs() < 1e-8 {
245 let finite_count = data.iter().filter(|v| v.is_finite()).count();
247 if !bins.is_empty() {
248 bins[0] = finite_count;
249 }
250 return bins;
251 }
252
253 let bin_size = (max - min) / NUM_BINS as f32;
254
255 for &value in data {
256 if !value.is_finite() {
257 continue;
258 }
259 let bin_idx = ((value - min) / bin_size).floor() as usize;
260 let bin_idx = bin_idx.min(NUM_BINS - 1);
261 bins[bin_idx] += 1;
262 }
263
264 bins
265}
266
267fn rebin(
269 old_bins: &[usize],
270 old_min: f32,
271 old_max: f32,
272 new_bins: &mut [usize],
273 new_min: f32,
274 new_max: f32,
275) {
276 if old_bins.is_empty() || new_bins.is_empty() {
277 return;
278 }
279 let old_range = old_max - old_min;
280 let new_range = new_max - new_min;
281 if old_range.abs() < 1e-8 || new_range.abs() < 1e-8 {
282 let total: usize = old_bins.iter().sum();
284 if total > 0 {
285 let center = (old_min + old_max) * 0.5;
286 let idx = ((center - new_min) / new_range * new_bins.len() as f32).floor() as usize;
287 let idx = idx.min(new_bins.len() - 1);
288 new_bins[idx] += total;
289 }
290 return;
291 }
292 let old_bin_size = old_range / old_bins.len() as f32;
293 let new_bin_count = new_bins.len();
294 for (i, &count) in old_bins.iter().enumerate() {
295 if count == 0 {
296 continue;
297 }
298 let center = old_min + (i as f32 + 0.5) * old_bin_size;
299 let new_idx = ((center - new_min) / new_range * new_bin_count as f32).floor() as usize;
300 let new_idx = new_idx.min(new_bin_count - 1);
301 new_bins[new_idx] += count;
302 }
303}
304
305pub fn calculate_optimal_range(data: &[f32], method: CalibrationMethod) -> (f32, f32) {
307 if data.is_empty() {
308 return (0.0, 0.0);
309 }
310
311 match method {
312 CalibrationMethod::MinMax => {
313 let min = data
314 .iter()
315 .copied()
316 .filter(|v| v.is_finite())
317 .fold(f32::INFINITY, f32::min);
318 let max = data
319 .iter()
320 .copied()
321 .filter(|v| v.is_finite())
322 .fold(f32::NEG_INFINITY, f32::max);
323 (min, max)
324 }
325
326 CalibrationMethod::Percentile(p) => {
327 let stats = ActivationStats::from_data(data);
328 let lower = stats.percentile(100.0 - p);
329 let upper = stats.percentile(p);
330 (lower, upper)
331 }
332
333 CalibrationMethod::Entropy => optimize_kl_divergence(data),
334
335 CalibrationMethod::MSE => optimize_mse(data),
336 }
337}
338
339pub fn calculate_optimal_range_from_stats(
347 stats: &ActivationStats,
348 method: CalibrationMethod,
349) -> (f32, f32) {
350 match method {
351 CalibrationMethod::MinMax => (stats.min(), stats.max()),
352
353 CalibrationMethod::Percentile(p) => {
354 let lower = stats.percentile(100.0 - p);
355 let upper = stats.percentile(p);
356 (lower, upper)
357 }
358
359 CalibrationMethod::Entropy => optimize_kl_from_stats(stats),
360
361 CalibrationMethod::MSE => optimize_mse_from_stats(stats),
362 }
363}
364
365fn optimize_kl_divergence(data: &[f32]) -> (f32, f32) {
367 let stats = ActivationStats::from_data(data);
368
369 let candidates = [99.0, 99.5, 99.9, 99.95, 99.99];
371 let mut best_range = (stats.min, stats.max);
372 let mut best_kl = f32::INFINITY;
373
374 for &percentile in &candidates {
375 let lower = stats.percentile(100.0 - percentile);
376 let upper = stats.percentile(percentile);
377
378 let kl = calculate_kl_divergence(data, lower, upper);
379
380 if kl < best_kl {
381 best_kl = kl;
382 best_range = (lower, upper);
383 }
384 }
385
386 best_range
387}
388
389fn optimize_mse(data: &[f32]) -> (f32, f32) {
391 let stats = ActivationStats::from_data(data);
392
393 let candidates = [99.0, 99.5, 99.9, 99.95, 99.99];
395 let mut best_range = (stats.min, stats.max);
396 let mut best_mse = f32::INFINITY;
397
398 for &percentile in &candidates {
399 let lower = stats.percentile(100.0 - percentile);
400 let upper = stats.percentile(percentile);
401
402 let mse = calculate_quantization_mse(data, lower, upper);
403
404 if mse < best_mse {
405 best_mse = mse;
406 best_range = (lower, upper);
407 }
408 }
409
410 best_range
411}
412
413fn calculate_kl_divergence(data: &[f32], min: f32, max: f32) -> f32 {
418 if (max - min).abs() < 1e-8 {
419 return 0.0;
420 }
421
422 let num_bins = 128;
423 let bin_size = (max - min) / num_bins as f32;
424 let scale = (max - min) / 255.0;
425
426 let mut orig_bins = vec![0usize; num_bins];
427 let mut quant_bins = vec![0usize; num_bins];
428
429 for &v in data {
430 let clipped = v.clamp(min, max);
431
432 let bin = ((clipped - min) / bin_size).floor() as usize;
434 let bin = bin.min(num_bins - 1);
435 orig_bins[bin] += 1;
436
437 let q = ((clipped - min) / scale).round();
439 let dequant = min + q * scale;
440 let qbin = ((dequant.clamp(min, max) - min) / bin_size).floor() as usize;
441 let qbin = qbin.min(num_bins - 1);
442 quant_bins[qbin] += 1;
443 }
444
445 let n = data.len() as f32;
446 let epsilon = 1e-10_f32;
447 let mut kl = 0.0_f32;
448
449 for i in 0..num_bins {
450 let p = (orig_bins[i] as f32 + epsilon) / (n + epsilon * num_bins as f32);
451 let q = (quant_bins[i] as f32 + epsilon) / (n + epsilon * num_bins as f32);
452 kl += p * (p / q).ln();
453 }
454
455 kl
456}
457
458fn calculate_quantization_mse(data: &[f32], min: f32, max: f32) -> f32 {
459 if (max - min).abs() < 1e-8 {
460 return 0.0;
461 }
462
463 let scale = (max - min) / 255.0;
464
465 let mse: f32 = data
466 .iter()
467 .map(|&v| {
468 let clipped = v.clamp(min, max);
469 let q = ((clipped - min) / scale).round().clamp(0.0, 255.0);
470 let dequantized = min + q * scale;
471 (v - dequantized).powi(2)
472 })
473 .sum::<f32>()
474 / data.len() as f32;
475
476 mse
477}
478
479fn histogram_kl_divergence(stats: &ActivationStats, min: f32, max: f32) -> f32 {
490 if (max - min).abs() < 1e-8 {
491 return 0.0;
492 }
493 let hist = stats.histogram_data();
494 if hist.is_empty() {
495 return 0.0;
496 }
497
498 const NUM_REBINS: usize = 128;
499 let rebin_size = (max - min) / NUM_REBINS as f32;
500 let scale = (max - min) / 255.0;
501
502 let mut orig = vec![0.0_f32; NUM_REBINS];
503 let mut quant = vec![0.0_f32; NUM_REBINS];
504
505 for &(center, count) in &hist {
506 let clipped = center.clamp(min, max);
507 let count_f = count as f32;
508
509 let bin = ((clipped - min) / rebin_size).floor() as usize;
510 let bin = bin.min(NUM_REBINS - 1);
511 orig[bin] += count_f;
512
513 let q = ((clipped - min) / scale).round();
514 let dq = min + q * scale;
515 let qbin = ((dq.clamp(min, max) - min) / rebin_size).floor() as usize;
516 let qbin = qbin.min(NUM_REBINS - 1);
517 quant[qbin] += count_f;
518 }
519
520 let total: f32 = orig.iter().sum();
521 if total == 0.0 {
522 return 0.0;
523 }
524
525 let epsilon = 1e-10_f32;
526 let denom = total + epsilon * NUM_REBINS as f32;
527 let mut kl = 0.0_f32;
528 for i in 0..NUM_REBINS {
529 let p = (orig[i] + epsilon) / denom;
530 let q = (quant[i] + epsilon) / denom;
531 kl += p * (p / q).ln();
532 }
533 kl
534}
535
536fn histogram_quantization_mse(stats: &ActivationStats, min: f32, max: f32) -> f32 {
539 if (max - min).abs() < 1e-8 {
540 return 0.0;
541 }
542
543 let scale = (max - min) / 255.0;
544 let mut weighted_sse = 0.0_f64;
545 let mut total_count = 0_u64;
546
547 for (center, count) in stats.histogram_data() {
548 let clipped = center.clamp(min, max);
549 let q = ((clipped - min) / scale).round().clamp(0.0, 255.0);
550 let dq = min + q * scale;
551 let err = (center - dq) as f64;
552 weighted_sse += err * err * count as f64;
553 total_count += count as u64;
554 }
555
556 if total_count == 0 {
557 0.0
558 } else {
559 (weighted_sse / total_count as f64) as f32
560 }
561}
562
563fn optimize_kl_from_stats(stats: &ActivationStats) -> (f32, f32) {
564 let candidates = [99.0, 99.5, 99.9, 99.95, 99.99];
565 let mut best_range = (stats.min(), stats.max());
566 let mut best_kl = f32::INFINITY;
567
568 for &percentile in &candidates {
569 let lower = stats.percentile(100.0 - percentile);
570 let upper = stats.percentile(percentile);
571 let kl = histogram_kl_divergence(stats, lower, upper);
572 if kl < best_kl {
573 best_kl = kl;
574 best_range = (lower, upper);
575 }
576 }
577 best_range
578}
579
580fn optimize_mse_from_stats(stats: &ActivationStats) -> (f32, f32) {
581 let candidates = [99.0, 99.5, 99.9, 99.95, 99.99];
582 let mut best_range = (stats.min(), stats.max());
583 let mut best_mse = f32::INFINITY;
584
585 for &percentile in &candidates {
586 let lower = stats.percentile(100.0 - percentile);
587 let upper = stats.percentile(percentile);
588 let mse = histogram_quantization_mse(stats, lower, upper);
589 if mse < best_mse {
590 best_mse = mse;
591 best_range = (lower, upper);
592 }
593 }
594 best_range
595}
596
597#[cfg(test)]
598mod tests {
599 use super::*;
600
601 #[test]
602 fn test_activation_stats() {
603 let data = vec![-1.0, -0.5, 0.0, 0.5, 1.0];
604 let stats = ActivationStats::from_data(&data);
605
606 assert_eq!(stats.min(), -1.0);
607 assert_eq!(stats.max(), 1.0);
608 assert!((stats.mean() - 0.0).abs() < 0.01);
609
610 let p50 = stats.percentile(50.0);
611 assert!((p50 - 0.0).abs() < 0.3);
612 }
613
614 #[test]
615 fn test_update_on_default_stats_with_zero_data_does_not_panic() {
616 let mut stats = ActivationStats::default();
621 stats.update(&[0.0, 0.0, 0.0]);
622 assert_eq!(stats.count(), 3);
623 assert_eq!(stats.min(), 0.0);
624 assert_eq!(stats.max(), 0.0);
625 assert_eq!(stats.percentile(50.0), 0.0);
627 }
628
629 #[test]
630 fn test_update_after_empty_from_data_bootstraps() {
631 let mut stats = ActivationStats::from_data(&[]);
634 assert_eq!(stats.count(), 0);
635 stats.update(&[1.0, 2.0, 3.0, 4.0]);
636 assert_eq!(stats.count(), 4);
637 assert_eq!(stats.min(), 1.0);
638 assert_eq!(stats.max(), 4.0);
639 stats.update(&[10.0]);
641 assert_eq!(stats.count(), 5);
642 assert_eq!(stats.max(), 10.0);
643 }
644
645 #[test]
646 fn test_update_after_all_nan_from_data_then_zero_data() {
647 let mut stats = ActivationStats::from_data(&[f32::NAN, f32::NAN]);
650 assert_eq!(stats.count(), 0);
651 stats.update(&[0.0, 0.0]);
652 assert_eq!(stats.count(), 2);
653 assert_eq!(stats.min(), 0.0);
654 assert_eq!(stats.max(), 0.0);
655 }
656
657 #[test]
662 fn test_minmax_from_stats_matches_raw_data() {
663 let data: Vec<f32> = (0..1000).map(|i| (i as f32 - 500.0) / 500.0).collect();
664 let stats = ActivationStats::from_data(&data);
665
666 let from_stats = calculate_optimal_range_from_stats(&stats, CalibrationMethod::MinMax);
667 let from_raw = calculate_optimal_range(&data, CalibrationMethod::MinMax);
668
669 assert_eq!(from_stats.0, from_raw.0);
671 assert_eq!(from_stats.1, from_raw.1);
672 }
673
674 #[test]
675 fn test_percentile_from_stats_is_deterministic() {
676 let data: Vec<f32> = (0..500).map(|i| (i as f32 - 250.0) / 100.0).collect();
679 let stats = ActivationStats::from_data(&data);
680
681 let r1 = calculate_optimal_range_from_stats(&stats, CalibrationMethod::Percentile(99.9));
682 let r2 = calculate_optimal_range_from_stats(&stats, CalibrationMethod::Percentile(99.9));
683 let r3 = calculate_optimal_range_from_stats(&stats, CalibrationMethod::Percentile(99.9));
684
685 assert_eq!(r1, r2);
686 assert_eq!(r2, r3);
687 }
688
689 #[test]
690 fn test_mse_from_stats_is_deterministic() {
691 let data: Vec<f32> = (0..500).map(|i| (i as f32 - 250.0) / 100.0).collect();
692 let stats = ActivationStats::from_data(&data);
693
694 let r1 = calculate_optimal_range_from_stats(&stats, CalibrationMethod::MSE);
695 let r2 = calculate_optimal_range_from_stats(&stats, CalibrationMethod::MSE);
696 assert_eq!(r1, r2);
697 }
698
699 #[test]
700 fn test_entropy_from_stats_is_deterministic() {
701 let data: Vec<f32> = (0..500).map(|i| (i as f32 - 250.0) / 100.0).collect();
702 let stats = ActivationStats::from_data(&data);
703
704 let r1 = calculate_optimal_range_from_stats(&stats, CalibrationMethod::Entropy);
705 let r2 = calculate_optimal_range_from_stats(&stats, CalibrationMethod::Entropy);
706 assert_eq!(r1, r2);
707 }
708
709 #[test]
710 fn test_all_methods_produce_finite_ranges() {
711 let data: Vec<f32> = (0..200).map(|i| (i as f32 / 50.0) - 1.0).collect();
714 let stats = ActivationStats::from_data(&data);
715
716 for method in [
717 CalibrationMethod::MinMax,
718 CalibrationMethod::Percentile(99.9),
719 CalibrationMethod::Entropy,
720 CalibrationMethod::MSE,
721 ] {
722 let (lo, hi) = calculate_optimal_range_from_stats(&stats, method);
723 assert!(lo.is_finite(), "{:?}: lower bound not finite", method);
724 assert!(hi.is_finite(), "{:?}: upper bound not finite", method);
725 assert!(lo <= hi, "{:?}: lo ({}) > hi ({})", method, lo, hi);
726 }
727 }
728
729 #[test]
730 fn test_stats_based_matches_raw_based_on_bulk_data() {
731 let data: Vec<f32> = (0..1000).map(|i| (i as f32 - 500.0) / 100.0).collect();
735 let stats = ActivationStats::from_data(&data);
736
737 let from_stats =
738 calculate_optimal_range_from_stats(&stats, CalibrationMethod::Percentile(99.0));
739 let from_raw = calculate_optimal_range(&data, CalibrationMethod::Percentile(99.0));
740
741 let width = stats.max() - stats.min();
742 let bin_width = width / 256.0;
743 let tolerance = 3.0 * bin_width + 1e-4;
744 assert!(
745 (from_stats.0 - from_raw.0).abs() <= tolerance,
746 "lower percentile drift: stats={} raw={} tol={}",
747 from_stats.0,
748 from_raw.0,
749 tolerance
750 );
751 assert!(
752 (from_stats.1 - from_raw.1).abs() <= tolerance,
753 "upper percentile drift: stats={} raw={} tol={}",
754 from_stats.1,
755 from_raw.1,
756 tolerance
757 );
758 }
759}