proof_stream/viz/
inference.rs

1//! Helpers for building `ProofEvent::LayerActivation` from intermediate tensors.
2
3use crate::events::{ActivationStats, LayerKind, ProofEvent};
4
5/// Build a `LayerActivation` event from a flat slice of M31 raw values (u32).
6///
7/// At most `max_sample` elements are included in `output_sample`.
8/// Stats (mean, std, min, max, sparsity) are computed over the full slice.
9pub fn layer_activation_event(
10    layer_idx: usize,
11    node_id: usize,
12    kind: LayerKind,
13    output_shape: (usize, usize),
14    values: &[u32],
15    max_sample: usize,
16) -> ProofEvent {
17    let stats = compute_stats(values);
18    let sample: Vec<u32> = values.iter().copied().take(max_sample).collect();
19
20    ProofEvent::LayerActivation {
21        layer_idx,
22        node_id,
23        kind,
24        output_shape,
25        output_sample: sample,
26        stats,
27    }
28}
29
30/// Compute activation statistics over raw M31 values.
31fn compute_stats(values: &[u32]) -> ActivationStats {
32    if values.is_empty() {
33        return ActivationStats {
34            mean: 0.0,
35            std_dev: 0.0,
36            min: 0.0,
37            max: 0.0,
38            sparsity: 1.0,
39        };
40    }
41
42    let p = 0x7fff_ffff_u32 as f64;
43    let n = values.len() as f64;
44
45    let sum: f64 = values.iter().map(|&v| v as f64 / p).sum();
46    let mean = (sum / n) as f32;
47
48    let var: f64 = values
49        .iter()
50        .map(|&v| {
51            let x = v as f64 / p - mean as f64;
52            x * x
53        })
54        .sum::<f64>()
55        / n;
56    let std_dev = var.sqrt() as f32;
57
58    let min = values.iter().map(|&v| v as f64 / p).fold(f64::INFINITY, f64::min) as f32;
59    let max = values.iter().map(|&v| v as f64 / p).fold(f64::NEG_INFINITY, f64::max) as f32;
60
61    let zeros = values.iter().filter(|&&v| v == 0).count();
62    let sparsity = zeros as f32 / values.len() as f32;
63
64    ActivationStats {
65        mean,
66        std_dev,
67        min,
68        max,
69        sparsity,
70    }
71}
proof_stream/viz/inference.rs

proof_stream/viz/
inference.rs
