| 26 | |
| 27 | template<typename T, int dim> |
| 28 | void sortBatched(Array<T>& val, bool isAscending) { |
| 29 | af::dim4 inDims = val.dims(); |
| 30 | |
| 31 | // Sort dimension |
| 32 | af::dim4 tileDims(1); |
| 33 | af::dim4 seqDims = inDims; |
| 34 | tileDims[dim] = inDims[dim]; |
| 35 | seqDims[dim] = 1; |
| 36 | |
| 37 | Array<uint> key = iota<uint>(seqDims, tileDims); |
| 38 | |
| 39 | Array<uint> resKey = createEmptyArray<uint>(dim4()); |
| 40 | Array<T> resVal = createEmptyArray<T>(dim4()); |
| 41 | |
| 42 | val.setDataDims(inDims.elements()); |
| 43 | key.setDataDims(inDims.elements()); |
| 44 | |
| 45 | sort_by_key<T, uint>(resVal, resKey, val, key, 0, isAscending); |
| 46 | |
| 47 | // Needs to be ascending (true) in order to maintain the indices properly |
| 48 | sort_by_key<uint, T>(key, val, resKey, resVal, 0, true); |
| 49 | val.setDataDims(inDims); // This is correct only for dim0 |
| 50 | } |
| 51 | |
| 52 | template<typename T> |
| 53 | void sort0(Array<T>& val, bool isAscending) { |
nothing calls this directly
no test coverage detected