Function pool_states 

pub fn pool_states(
    k: &[f32],
    gate: &[f32],
    valid: &[u8],
    ape: &[f32],
    dims: DsaDims,
    seq: usize,
) -> Pools

2026-09-25: Build the pools. k/gate are [seq, index_head_dim], valid is [seq]. The indexer golden records the output of transformers’ get_pooled_states for comparison.