Function exec_set

dnn/src/naive/indexing_one_hot/opr_impl.cpp:47–73

template <typename data_type, typename idx_type = dt_int32>
void exec_set(
        const TensorND& data, const TensorND& index, const TensorND& sub,
        uint32_t axis) {
    TensorND data_nomid = data;
    data_nomid.layout.remove_axis_inplace(axis);
    auto data_mid_stride = data.layout.stride[axis];
    int data_mid_shape = data.layout.shape[axis];

    size_t nr_elems = data_nomid.layout.total_nr_elems();
    megdnn_assert(
            nr_elems == index.layout.total_nr_elems() &&
            nr_elems == sub.layout.total_nr_elems());
    auto data_iter = tensor_iter_valonly<data_type>(data_nomid).begin();
    auto idx_iter = tensor_iter_valonly<idx_type>(index).begin();
    auto sub_iter = tensor_iter_valonly<data_type>(sub).begin();

    data_type* dptr = data.ptr<data_type>();

    for (size_t i = 0; i < nr_elems; ++i) {
        auto idx = *idx_iter;
        megdnn_assert(idx >= 0 && idx < data_mid_shape);
        dptr[data_iter.offset() + *idx_iter * data_mid_stride] = *sub_iter;
        ++data_iter;
        ++sub_iter;
        ++idx_iter;
    }
}

}  // anonymous namespace