| 12 | |
| 13 | template <typename data_type, typename idx_type = dt_int32> |
| 14 | void exec_get( |
| 15 | const TensorND& src, const TensorND& index, const TensorND& dst, |
| 16 | uint32_t axis) { |
| 17 | TensorND src_nomid = src; |
| 18 | src_nomid.layout.remove_axis_inplace(axis); |
| 19 | auto src_mid_stride = src.layout.stride[axis]; |
| 20 | int src_mid_shape = src.layout.shape[axis]; |
| 21 | |
| 22 | size_t nr_elems = src_nomid.layout.total_nr_elems(); |
| 23 | megdnn_assert( |
| 24 | nr_elems == index.layout.total_nr_elems() && |
| 25 | nr_elems == dst.layout.total_nr_elems()); |
| 26 | auto src_iter = tensor_iter_valonly<data_type>(src_nomid).begin(); |
| 27 | auto idx_iter = tensor_iter_valonly<idx_type>(index).begin(); |
| 28 | auto dst_iter = tensor_iter_valonly<data_type>(dst).begin(); |
| 29 | |
| 30 | data_type* sptr = src.ptr<data_type>(); |
| 31 | |
| 32 | for (size_t i = 0; i < nr_elems; ++i) { |
| 33 | auto idx = *idx_iter; |
| 34 | megdnn_assert( |
| 35 | idx >= 0 && idx < src_mid_shape, |
| 36 | "bad value in IndexingOneHot index: input shape is %d, " |
| 37 | "index value is %d", |
| 38 | src_mid_shape, idx); |
| 39 | *dst_iter = sptr[src_iter.offset() + *idx_iter * src_mid_stride]; |
| 40 | ++src_iter; |
| 41 | ++dst_iter; |
| 42 | ++idx_iter; |
| 43 | } |
| 44 | } |
| 45 | |
| 46 | template <typename data_type, typename idx_type = dt_int32> |
| 47 | void exec_set( |
nothing calls this directly
no test coverage detected