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