| 62 | |
| 63 | template<typename T> |
| 64 | Array<T> sort(const Array<T>& in, const unsigned dim, bool isAscending) { |
| 65 | Array<T> out = copyArray<T>(in); |
| 66 | switch (dim) { |
| 67 | case 0: sort0<T>(out, isAscending); break; |
| 68 | case 1: sortBatched<T, 1>(out, isAscending); break; |
| 69 | case 2: sortBatched<T, 2>(out, isAscending); break; |
| 70 | case 3: sortBatched<T, 3>(out, isAscending); break; |
| 71 | default: AF_ERROR("Not Supported", AF_ERR_NOT_SUPPORTED); |
| 72 | } |
| 73 | |
| 74 | if (dim != 0) { |
| 75 | af::dim4 preorderDims = out.dims(); |
| 76 | af::dim4 reorderDims(0, 1, 2, 3); |
| 77 | reorderDims[dim] = 0; |
| 78 | preorderDims[0] = out.dims()[dim]; |
| 79 | for (int i = 1; i <= static_cast<int>(dim); i++) { |
| 80 | reorderDims[i - 1] = i; |
| 81 | preorderDims[i] = out.dims()[i - 1]; |
| 82 | } |
| 83 | |
| 84 | out.setDataDims(preorderDims); |
| 85 | out = reorder<T>(out, reorderDims); |
| 86 | } |
| 87 | return out; |
| 88 | } |
| 89 | |
| 90 | #define INSTANTIATE(T) \ |
| 91 | template Array<T> sort<T>(const Array<T>& in, const unsigned dim, \ |
no test coverage detected