index_scores

Function index_scores 

Source
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.