Struct ProbePool 

pub struct ProbePool<'a> { /* private fields */ }

Pools a probe head’s first softmax as positions arrive, so the cells are never all held at once.

Materialising them costs seq * cells * d_model floats — 504 MB at full context on the shipped model, which is not survivable in a browser tab. This holds cells * probes * d_model regardless of sequence length: 252 KB.

The trick is the online softmax. r = Σ_t softmax(s)_t · cell_t is computed by carrying a running maximum, denominator and numerator, and rescaling both whenever a larger score appears — algebraically identical to the two-pass form, which ProbeHead::forward still does and the tests compare against.

Implementations

impl<'a> ProbePool<'a>

pub fn new(head: &'a ProbeHead, d_model: usize) -> Self

pub fn observe(&mut self, cell_index: usize, cell: &[f32])

Feed one position’s cell for one cell index (0 is the input embedding).

Order does not matter — the pooling is a sum over positions — so a caller may stream all positions for one cell before moving to the next, which is how the forward pass is laid out.

pub fn finish(self) -> Vec<f32>

Finish the pooling and run the head’s second softmax and projection.