| 27 | } |
| 28 | |
| 29 | void expansion_index( |
| 30 | const dt_int32* const index_arr, const size_t index_size, const TensorND* rst, |
| 31 | const size_t* shape_arr, const int ndim) { |
| 32 | SmallVector<int, 8> shape_reverse_multiply_reduce_arr({1}); |
| 33 | for (int div_index = 1; div_index < ndim; div_index++) { |
| 34 | shape_reverse_multiply_reduce_arr[div_index] = |
| 35 | shape_arr[ndim - div_index] * |
| 36 | shape_reverse_multiply_reduce_arr[div_index - 1]; |
| 37 | } |
| 38 | |
| 39 | for (int dim_pos = 0; dim_pos < ndim; dim_pos++) { |
| 40 | dt_int32* dim_pt = rst->ptr<dt_int32>() + index_size * dim_pos; |
| 41 | for (size_t ele_pos = 0; ele_pos < index_size; ele_pos++) { |
| 42 | int dim_pos_of_ele = index_arr[ele_pos] / |
| 43 | shape_reverse_multiply_reduce_arr[ndim - 1 - dim_pos]; |
| 44 | int dim_index_of_ele = dim_pos_of_ele % shape_arr[dim_pos]; |
| 45 | dim_pt[ele_pos] = dim_index_of_ele; |
| 46 | } |
| 47 | } |
| 48 | } |
| 49 | |
| 50 | TensorND NonZeroImpl::exec( |
| 51 | _megdnn_tensor_in src, _megdnn_workspace workspace, |