| 38 | } // namespace |
| 39 | |
| 40 | af_err af_topk(af_array *values, af_array *indices, const af_array in, |
| 41 | const int k, const int dim, const af_topk_function order) { |
| 42 | try { |
| 43 | af::topkFunction ord = (order == AF_TOPK_DEFAULT ? AF_TOPK_MAX : order); |
| 44 | |
| 45 | const ArrayInfo &inInfo = getInfo(in); |
| 46 | |
| 47 | ARG_ASSERT(2, (inInfo.ndims() > 0)); |
| 48 | |
| 49 | if (inInfo.elements() == 1) { |
| 50 | dim_t dims[1] = {1}; |
| 51 | af_err errValue = af_constant(indices, 0, 1, dims, u32); |
| 52 | return errValue == AF_SUCCESS ? af_retain_array(values, in) |
| 53 | : errValue; |
| 54 | } |
| 55 | |
| 56 | int rdim = dim; |
| 57 | const auto &inDims = inInfo.dims(); |
| 58 | |
| 59 | if (rdim == -1) { |
| 60 | for (dim_t d = 0; d < 4; d++) { |
| 61 | if (inDims[d] > 1) { |
| 62 | rdim = d; |
| 63 | break; |
| 64 | } |
| 65 | } |
| 66 | } |
| 67 | |
| 68 | ARG_ASSERT(2, (inInfo.dims()[rdim] >= k)); |
| 69 | ARG_ASSERT( |
| 70 | 4, (k > 0) && (k <= 256)); // TODO(umar): Remove this limitation |
| 71 | |
| 72 | if (rdim != 0) { |
| 73 | AF_ERROR("topk is supported along dimenion 0 only.", |
| 74 | AF_ERR_NOT_SUPPORTED); |
| 75 | } |
| 76 | |
| 77 | af_dtype type = inInfo.getType(); |
| 78 | |
| 79 | switch (type) { |
| 80 | // TODO(umar): FIX RETURN VALUES HERE |
| 81 | case f32: topk<float>(values, indices, in, k, rdim, ord); break; |
| 82 | case f64: topk<double>(values, indices, in, k, rdim, ord); break; |
| 83 | case u32: topk<uint>(values, indices, in, k, rdim, ord); break; |
| 84 | case s32: topk<int>(values, indices, in, k, rdim, ord); break; |
| 85 | case f16: topk<half>(values, indices, in, k, rdim, ord); break; |
| 86 | default: TYPE_ERROR(1, type); |
| 87 | } |
| 88 | } |
| 89 | CATCHALL; |
| 90 | |
| 91 | return AF_SUCCESS; |
| 92 | } |
no test coverage detected