def allgather_sizes(send_data, world_size, num_parts, return_sizes=False):
    """
    Perform all gather on list lengths, used to compute prefix sums
    to determine the offsets on each ranks. This is used to allocate
    global ids for edges/nodes on each ranks.

    Parameters
    ----------
    send_data : numpy array
        Data on which allgather is performed.
    world_size : integer
        No. of processes configured for execution
    num_parts : integer
        No. of output graph partitions
    return_sizes : bool
        Boolean flag to indicate whether to return raw sizes from each process
        or perform prefix sum on the raw sizes.

    Returns :
    ---------
        numpy array
            array with the prefix sum
    """

    # Assert on the world_size, num_parts
    assert (num_parts % world_size) == 0

    # compute the length of the local data
    send_length = len(send_data)
    out_tensor = torch.as_tensor(send_data, dtype=torch.int64)
    in_tensor = [
        torch.zeros(send_length, dtype=torch.int64) for _ in range(world_size)
    ]

    # all_gather message
    dist.all_gather(in_tensor, out_tensor)

    # Return on the raw sizes from each process
    if return_sizes:
        return torch.cat(in_tensor).numpy()

    # gather sizes in on array to return to the invoking function
    rank_sizes = np.zeros(num_parts + 1, dtype=np.int64)
    part_counts = torch.cat(in_tensor).numpy()

    count = rank_sizes[0]
    idx = 1
    for local_part_id in range(num_parts // world_size):
        for r in range(world_size):
            count += part_counts[r * (num_parts // world_size) + local_part_id]
            rank_sizes[idx] = count
            idx += 1

    return rank_sizes

    Perform all gather on list lengths, used to compute prefix sums
    to determine the offsets on each ranks. This is used to allocate
    global ids for edges/nodes on each ranks.

    Parameters
    ----------
    send_data : numpy array
        Data on which allgather is performed.
    world_size : integer
        No. of processes configured for execution
    num_parts : integer
        No. of output graph partitions
    return_sizes : bool
        Boolean flag to indicate whether to return raw sizes from each process
        or perform prefix sum on the raw sizes.

    Returns :
    ---------
        numpy array
            array with the prefix sum
    
tools/distpartitioning/gloo_wrapper.py