Top-K selection: returns `(values, indices)` for the largest K elements along the last dimension. Uses partial sort O(N + K log K) for large N.
(&self, x: &Tensor, k: usize)
| 555 | /// Top-K selection: returns `(values, indices)` for the largest K elements |
| 556 | /// along the last dimension. Uses partial sort O(N + K log K) for large N. |
| 557 | fn topk(&self, x: &Tensor, k: usize) -> Result<(Tensor, Tensor)> { |
| 558 | // Ensure F32 — sort/partial_sort and to_vec2::<f32> require F32 input |
| 559 | let x = if x.dtype() != candle_core::DType::F32 { |
| 560 | x.to_dtype(candle_core::DType::F32)?.contiguous()? |
| 561 | } else { |
| 562 | x.contiguous()? |
| 563 | }; |
| 564 | let shape = x.shape(); |
| 565 | let last = shape.dims().last().copied().unwrap_or(0); |
| 566 | if last <= 32 || k * 2 >= last { |
| 567 | let last_dim = x.rank() - 1; |
| 568 | let (sorted, indices) = x.sort_last_dim(false)?; |
| 569 | let top_vals = sorted.narrow(last_dim, 0, k)?; |
| 570 | let top_idx = indices.narrow(last_dim, 0, k)?; |
| 571 | return Ok((top_vals, top_idx)); |
| 572 | } |
| 573 | let rank = x.rank(); |
| 574 | let batch: usize = shape.dims()[..rank - 1].iter().product(); |
| 575 | let flat = x.reshape((batch, last))?; |
| 576 | let data = flat.to_vec2::<f32>()?; |
| 577 | let mut all_vals = Vec::with_capacity(batch * k); |
| 578 | let mut all_idxs = Vec::with_capacity(batch * k); |
| 579 | for row in &data { |
| 580 | let mut indices: Vec<u32> = (0..last as u32).collect(); |
| 581 | indices.select_nth_unstable_by(k - 1, |&a, &b| { |
| 582 | row[b as usize] |
| 583 | .partial_cmp(&row[a as usize]) |
| 584 | .unwrap_or(std::cmp::Ordering::Equal) |
| 585 | }); |
| 586 | let top_slice = &mut indices[..k]; |
| 587 | top_slice.sort_unstable_by(|&a, &b| { |
| 588 | row[b as usize] |
| 589 | .partial_cmp(&row[a as usize]) |
| 590 | .unwrap_or(std::cmp::Ordering::Equal) |
| 591 | }); |
| 592 | for &idx in top_slice.iter() { |
| 593 | all_vals.push(row[idx as usize]); |
| 594 | all_idxs.push(idx); |
| 595 | } |
| 596 | } |
| 597 | let mut out_shape: Vec<usize> = shape.dims()[..rank - 1].to_vec(); |
| 598 | out_shape.push(k); |
| 599 | let vals = Tensor::from_vec(all_vals, out_shape.as_slice(), x.device())?; |
| 600 | let idxs = Tensor::from_vec(all_idxs, out_shape.as_slice(), x.device())?; |
| 601 | Ok((vals, idxs)) |
| 602 | } |
| 603 | |
| 604 | // ── Convolutions ────────────────────────────────────────────────── |
| 605 |