| 38 | namespace arrayfire { |
| 39 | namespace opencl { |
| 40 | vector<af_index_t> indexForTopK(const int k) { |
| 41 | af_index_t idx; |
| 42 | idx.idx.seq = af_seq{0.0, static_cast<double>(k) - 1.0, 1.0}; |
| 43 | idx.isSeq = true; |
| 44 | idx.isBatch = false; |
| 45 | |
| 46 | af_index_t sp; |
| 47 | sp.idx.seq = af_span; |
| 48 | sp.isSeq = true; |
| 49 | sp.isBatch = false; |
| 50 | |
| 51 | return vector<af_index_t>({idx, sp, sp, sp}); |
| 52 | } |
| 53 | |
| 54 | template<typename T> |
| 55 | void topk(Array<T>& vals, Array<unsigned>& idxs, const Array<T>& in, |