| 34 | |
| 35 | template<typename T> |
| 36 | static inline af_array reorder(const af_array in, const af::dim4 &rdims0) { |
| 37 | Array<T> In = detail::createEmptyArray<T>(af::dim4(0)); |
| 38 | dim4 rdims = rdims0; |
| 39 | |
| 40 | if (rdims[0] == 1 && rdims[1] == 0) { |
| 41 | In = transpose(getArray<T>(in), false); |
| 42 | std::swap(rdims[0], rdims[1]); |
| 43 | } else { |
| 44 | In = getArray<T>(in); |
| 45 | } |
| 46 | const dim4 idims = In.dims(); |
| 47 | const dim4 istrides = In.strides(); |
| 48 | |
| 49 | // Ensure all JIT nodes are evaled |
| 50 | In.eval(); |
| 51 | |
| 52 | af_array out; |
| 53 | if (rdims[0] == 0 && rdims[1] == 1 && rdims[2] == 2 && rdims[3] == 3) { |
| 54 | out = getHandle(In); |
| 55 | } else if (rdims[0] == 0) { |
| 56 | dim4 odims = dim4(1, 1, 1, 1); |
| 57 | dim4 ostrides = dim4(1, 1, 1, 1); |
| 58 | for (int i = 0; i < 4; i++) { |
| 59 | odims[i] = idims[rdims[i]]; |
| 60 | ostrides[i] = istrides[rdims[i]]; |
| 61 | } |
| 62 | Array<T> Out = In; |
| 63 | // Use modDims instead of setDataDims to only modify the ArrayInfo |
| 64 | Out.modDims(odims); |
| 65 | Out.modStrides(ostrides); |
| 66 | out = getHandle(Out); |
| 67 | } else { |
| 68 | Array<T> Out = reorder<T>(In, rdims); |
| 69 | out = getHandle(Out); |
| 70 | } |
| 71 | return out; |
| 72 | } |
| 73 | |
| 74 | af_err af_reorder(af_array *out, const af_array in, const af::dim4 &rdims) { |
| 75 | try { |