PyTorch/ATen · CPU kthvalue

选择张量切片的第 k 小值及其位置

PyTorch 的 CPU kthvalue 使用 Quickselect，只继续处理包含目标秩的分区，并同步交换原始索引。

v2.5.1检索与计算核验 2026-10-12

工程问题#

张量计算有时只需要某一维的第 k 小值及其索引。对每条切片完整排序会计算不需要的其他顺序，还必须保持原输入不变并正确返回原始位置。

实际实现#

kthvalue_out_impl_cpu 克隆输入并分配索引缓冲。每条切片初始化原始索引，然后调用 quick_select_template：三数取中选择枢轴，双向扫描分区，按目标 k 缩小有效区间。交换回调同时移动值和索引，最终从 k-1 位置取结果。

处理流程#

验证 k 的取值范围，克隆输入并准备原始索引
逐切片选取枢轴并执行分区
依据枢轴位置只保留包含 k-1 的有效区间
每次交换同时更新值与索引，最后输出目标位置的数据

源码片段#

固定版本 · 原文摘录

aten/src/ATen/native/Sorting.cppv2.5.1 · L425–L442

查看原文件

425  TORCH_CHECK(k >= 1 && k <= slicesize,426              "kthvalue(): selected number k out of range for dimension ", dim);427428  at::assert_no_overlap(self, values);429430  _reduction_with_indices_allocate_or_resize_output(431      values, indices, self, dim_, keepdim);432  if (self.dim() == 0 && self.numel() == 1) {433    values.copy_(self);434    indices.zero_();435    return std::forward_as_tuple(values, indices);436  }437  auto tmp_values = self.clone(at::MemoryFormat::Contiguous);438  auto tmp_indices = at::empty(self.sizes(), self.options().dtype(kLong));439440  auto tmp_values_stride = tmp_values.strides()[dim];441  auto tmp_indices_stride = tmp_indices.strides()[dim];442  auto sizes = self.sizes();

kthvalue_out_impl_cpu：输入与工作缓冲

检查目标秩，对非标量常规路径克隆输入，并分配同尺寸的原始索引工作区。

CPPBSD-3-Clause（PyTorch，见同版 LICENSE）

aten/src/ATen/native/Sorting.cppv2.5.1 · L468–L486

查看原文件

468        for (const auto j : c10::irange(tmp_indices.size(0))) {469          tmp_indices[j] = j;470        }471472        // we want NaN to be sorted as top for numpy compatibility473        quick_select_template(474          tmp_values,475          k - 1,476          [](scalar_t x, scalar_t y) -> bool {477            return (478              (_isnan<scalar_t>(x) && !_isnan<scalar_t>(y)) || (x > y));479          },480          [&](int64_t i, int64_t j) {481            std::swap(tmp_values[i], tmp_values[j]);482            std::swap(tmp_indices[i], tmp_indices[j]);483          });484        *mode_value = tmp_values[k - 1];485        *mode_index = tmp_indices[k - 1];486      }

kthvalue_out_impl_cpu → quick_select_template

传入零基目标秩；比较器处理 NaN，交换回调同时重排值和原始索引，最后直接读取 k-1。

CPPBSD-3-Clause（PyTorch，见同版 LICENSE）

aten/src/ATen/native/Sorting.cppv2.5.1 · L148–L159

查看原文件

148    // Use median of three for pivot choice149    P = L + (R - L) / 2;150    swap_fn(P, L + 1);151    if (gt_or_nan(arr[L + 1], arr[R])) {152      swap_fn(L + 1, R);153    }154    if (gt_or_nan(arr[L], arr[R])) {155      swap_fn(L, R);156    }157    if (gt_or_nan(arr[L + 1], arr[L])) {158      swap_fn(L + 1, L);159    }

quick_select_template：三数取中

确定性三数取中用于选择枢轴，而不是随机枢轴或保证线性最坏界的 median-of-medians。

CPPBSD-3-Clause（PyTorch，见同版 LICENSE）

aten/src/ATen/native/Sorting.cppv2.5.1 · L161–L182

查看原文件

161    i = L + 1;162    j = R;163    piv = arr[L];164    do {165      do166        i++;167      while (gt_or_nan(piv, arr[i]));168      do169        j--;170      while (gt_or_nan(arr[j], piv));171      if (j < i)172        break;173      swap_fn(i, j);174    } while (true);175    swap_fn(L, j);176177    // Re-set active partition178    if (j <= k)179      L = i;180    if (j >= k)181      R = j - 1;182  } while (true);