topk_pools

Function topk_pools 

Source
pub fn topk_pools(
    scores: &[f32],
    valid_candidates: &[u8],
    n_pools: usize,
    q_rows: usize,
    select_k: usize,
) -> Vec<i32>
Expand description

Select up to select_k pools per query.

🔴 Deterministic tiebreak: higher score first, then SMALLER pool index. The reference uses torch.topk, whose tie order is implementation-defined, so the reference’s own pool identities are not a legal target on a tied row — only the selected set, and only when the tie does not straddle the cutoff. This function pins a total order so Atlas is reproducible regardless.