torch_cpu_topk_set

Function torch_cpu_topk_set 

Source
pub fn torch_cpu_topk_set(vals: &[f32], k: usize) -> Vec<usize>
Expand description

The index SET torch’s CPU topk(k, largest=True) returns for vals, ties included. For the row lengths here (k * 64 > n) torch runs std::nth_element on (value, index) pairs with a descending comparator and returns the first k pairs, so among equal values (the indexer’s many -inf entries) the picks are whatever libstdc++’s introselect leaves in front. That order is deterministic and this is a line-for-line emulation of it (__introselect, __move_median_to_first, __unguarded_partition, __insertion_sort). Pinned by topk_tie_order_matches_torch_cpu against this box’s torch.