def _broadcast_shape(
    data, rank, world_size, num_parts, is_feat_data, feat_name
):
    """Auxiliary function to broadcast the shape of a feature data.
    This information is used to figure out the type-ids for the
    local features.

    Parameters:
    -----------
    data : numpy array
        which is the feature data read from the disk
    rank : integer
        which represents the id of the process in the process group
    world_size : integer
        represents the total no. of process in the process group
    num_parts : integer
        specifying the no. of partitions
    is_feat_data : bool
        flag used to seperate feature data and edge data
    feat_name : string
        name of the feature

    Returns:
    -------
    list of tuples :
        which represents the range of type-ids for the data array.
    """
    assert len(data.shape) in [
        1,
        2,
    ], f"Data is expected to be 1-D or 2-D but got {data.shape}."
    data_shape = list(data.shape)

    if len(data_shape) == 1:
        data_shape.append(1)

    if is_feat_data:
        data_shape.append(DATA_TYPE_ID[data.dtype])

    data_shape = torch.tensor(data_shape, dtype=torch.int64)
    data_shape_output = [
        torch.zeros_like(data_shape) for _ in range(world_size)
    ]
    dist.all_gather(data_shape_output, data_shape)
    logging.debug(
        f"[Rank: {rank} Received shapes from all ranks: {data_shape_output}"
    )
    shapes = [x.numpy() for x in data_shape_output if x[0] != 0]
    shapes = np.vstack(shapes)

    if is_feat_data:
        logging.debug(
            f"shapes: {shapes}, condition: {all(shapes[0,2] == s for s in shapes[:,2])}"
        )
        assert all(
            shapes[0, 2] == s for s in shapes[:, 2]
        ), f"dtypes for {feat_name} does not match on all ranks"

    # compute tids here.
    type_counts = list(shapes[:, 0])
    tid_start = np.cumsum([0] + type_counts[:-1])
    tid_end = np.cumsum(type_counts)
    tid_ranges = list(zip(tid_start, tid_end))
    logging.debug(f"starts -> {tid_start} ... end -> {tid_end}")

    return tid_ranges
Auxiliary function to broadcast the shape of a feature data.
    This information is used to figure out the type-ids for the
    local features.

    Parameters:
    -----------
    data : numpy array
        which is the feature data read from the disk
    rank : integer
        which represents the id of the process in the process group
    world_size : integer
        represents the total no. of process in the process group
    num_parts : integer
        specifying the no. of partitions
    is_feat_data : bool
        flag used to seperate feature data and edge data
    feat_name : string
        name of the feature

    Returns:
    -------
    list of tuples :
        which represents the range of type-ids for the data array.
    
tools/distpartitioning/dataset_utils.py