Function topk

src/backend/oneapi/topk.cpp:48–56

template<typename T>
void topk(Array<T>& vals, Array<unsigned>& idxs, const Array<T>& in,
          const int k, const int dim, const af::topkFunction order) {
    auto values  = createEmptyArray<T>(in.dims());
    auto indices = createEmptyArray<unsigned>(in.dims());
    sort_index(values, indices, in, dim, order & AF_TOPK_MIN);
    auto indVec = indexForTopK(k);
    vals        = index<T>(values, indVec.data());
    idxs        = index<unsigned>(indices, indVec.data());
}

#define INSTANTIATE(T)                                                  \
    template void topk<T>(Array<T>&, Array<unsigned>&, const Array<T>&, \