| 23 | |
| 24 | template<typename T> |
| 25 | Array<T> join(const int dim, const Array<T> &first, const Array<T> &second) { |
| 26 | // All dimensions except join dimension must be equal |
| 27 | // Compute output dims |
| 28 | af::dim4 odims; |
| 29 | af::dim4 fdims = first.dims(); |
| 30 | af::dim4 sdims = second.dims(); |
| 31 | |
| 32 | for (int i = 0; i < 4; i++) { |
| 33 | if (i == dim) { |
| 34 | odims[i] = fdims[i] + sdims[i]; |
| 35 | } else { |
| 36 | odims[i] = fdims[i]; |
| 37 | } |
| 38 | } |
| 39 | |
| 40 | Array<T> out = createEmptyArray<T>(odims); |
| 41 | std::vector<CParam<T>> v{first, second}; |
| 42 | getQueue().enqueue(kernel::join<T>, dim, out, v, 2); |
| 43 | |
| 44 | return out; |
| 45 | } |
| 46 | |
| 47 | template<typename T> |
| 48 | void join(Array<T> &out, const int dim, const std::vector<Array<T>> &inputs) { |