xref: /aosp_15_r20/external/pytorch/aten/src/ATen/native/cuda/TensorTopK.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #pragma once
2 #include <cstdint>
3 
4 namespace at {
5 class TensorBase;
6 }
7 
8 namespace at {
9 namespace native {
10 void launch_gather_topk_kernel(
11     const TensorBase& self,
12     int64_t k, int64_t dim, bool largest,
13     const TensorBase& values, const TensorBase& indices);
14 }}
15