Function uniform_quantize

ops/cpp_modules/src/cpp_modules.cpp:288–334

py::array_t<int> uniform_quantize(py::array_t<int> seg_idx, py::array_t<float> residual, float acc) {
    // input array
    py::buffer_info si_buf = seg_idx.request();
    int *si_ptr = (int *) si_buf.ptr;
    py::buffer_info ri_buf = residual.request();
    float *ri_ptr = (float *) ri_buf.ptr;
    int h = si_buf.shape[0];
    int w = si_buf.shape[1];

    int cluster_num = 0;
    for (int h_i = 0; h_i < h; h_i++)
        for (int w_i = 0; w_i < w; w_i++)
            if (si_ptr[h_i * w + w_i] > cluster_num)
                cluster_num = si_ptr[h_i * w + w_i];
    cluster_num += 1;

    vector< vector<int> > quantized_residual;
    for (int i = 0; i < cluster_num; i++){
        vector<int> cur_qr;
        quantized_residual.push_back(cur_qr);
    }

    int total_num = 0;
    for (int h_i = 0; h_i < h; h_i++) {
        for (int w_i = 0; w_i < w; w_i++){
            int idx = si_ptr[h_i * w + w_i];
            if (idx != 1){
                quantized_residual[idx].push_back(round(ri_ptr[h_i * w + w_i] / acc));
                total_num++;
            }
        }
    }

    // return array
    auto quantized_residual_collect = py::array_t<int>(total_num);
    py::buffer_info qr_buf = quantized_residual_collect.request();
    int *qr_ptr = (int *) qr_buf.ptr;

    int k = 0;
    for (int i = 0; i < cluster_num; i++)
        for (std::vector<int>::iterator it = quantized_residual[i].begin() ; it != quantized_residual[i].end(); it++){
            qr_ptr[k] = *it;
            k++;
        }

    return quantized_residual_collect;
}


std::tuple<py::array_t<int>, py::array_t<int>> nonuniform_quantize(py::array_t<int> seg_idx,