pub fn index_scores(
q: &[f32],
weights: &[f32],
pools: &Pools,
dims: DsaDims,
q_rows: usize,
) -> Vec<f32>Expand description
Per-(query, pool) index score. Mirrors the matmul → relu → head-weighted sum chain.
q is [q_rows, index_heads, index_head_dim], weights is [q_rows, index_heads] and is
expected to ALREADY carry the index_heads^-0.5 factor.