Function pool_states 

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

Build the pools. k/gate are [seq, index_head_dim], valid is [seq].

Mirrors Glm5NextTextIndexer.get_pooled_states.