Function gather_points_v1

pointersect/pr/cuda/pr_cuda.cpp:222–235

Gather the points belonging to each grid cell. Args: grid_idxs: (b, n), grid linear index of each point total_cells: (b,) total grid cells valid_mask: (b, n), whether to record the point cell_counts: (b, max_num_cells), number of points in a cell Returns: vector of vector of tensors, b -> cell_idx -> pidx of points in the cell

//     vector of vector of tensors, b -> cell_idx -> pidx of points in the cell
// 
std::vector<std::vector<torch::Tensor> > gather_points_v1(
        const torch::Tensor & grid_idxs,  // (b, n), long
        const torch::Tensor & total_cells, // (b,) int64_t
        const torch::Tensor & valid_mask,  // (b, n), bool
        const torch::Tensor & cell_counts  // (b, n_cells), int32
) { 

    CHECK_INPUT(grid_idxs);
    CHECK_INPUT(total_cells);
    CHECK_INPUT(valid_mask);
    CHECK_INPUT(cell_counts);

    return gather_points_cuda(grid_idxs, total_cells, valid_mask, cell_counts);
}


// Gather the points belonging to each grid cell.