| 20 | namespace cuda { |
| 21 | template<typename in_t, typename idx_t> |
| 22 | Array<in_t> lookup(const Array<in_t> &input, const Array<idx_t> &indices, |
| 23 | const unsigned dim) { |
| 24 | const dim4 &iDims = input.dims(); |
| 25 | |
| 26 | dim4 oDims(1); |
| 27 | for (dim_t d = 0; d < 4; ++d) { |
| 28 | oDims[d] = (d == dim ? indices.elements() : iDims[d]); |
| 29 | } |
| 30 | |
| 31 | Array<in_t> out = createEmptyArray<in_t>(oDims); |
| 32 | |
| 33 | dim_t nDims = iDims.ndims(); |
| 34 | |
| 35 | kernel::lookup<in_t, idx_t>(out, input, indices, nDims, dim); |
| 36 | |
| 37 | return out; |
| 38 | } |
| 39 | |
| 40 | #define INSTANTIATE(T) \ |
| 41 | template Array<T> lookup<T, float>(const Array<T> &, const Array<float> &, \ |