[FEA] Expose algorithm selection in matrix.select_k
python API
#2427
Labels
feature request
New feature or request
matrix.select_k
python API
#2427
Is your feature request related to a problem? Please describe.
The current python API (here) in
pylibraft.matrix.select_k
does not expose the algorithm selection that the C++ API does, see here. Instead it uses the default value which automatically selects the algorithm.Describe the solution you'd like
I'd suggest passing an extra enum argument to the function that allows selecting the algorithm similar to the C++ API.
Describe alternatives you've considered
In a benchmark, the default automatic algorithm performs poorly (~10X slower) than the AIR-TopK reported in this paper. Using the C++ API, I can reproduce the same timings reported in the paper when choosing
raft::matrix::SelectAlgo::kRadix11bits
algorithm.Additional context
C++ documentation is provided here.
The text was updated successfully, but these errors were encountered: