Function batch_top_k_argmax 

pub fn batch_top_k_argmax(
    scores: &ArrayView2<'_, f32>,
    k: usize,
    parallel: bool,
) -> Result<Vec<Vec<(usize, f32)>>>

Batch top-k argmax over an (n_rows, n_cols) matrix. With parallel = true distributes rows across rayon’s pool.