def exchange_feature(
    rank,
    data,
    id_lookup,
    feat_type,
    feat_key,
    featdata_key,
    gid_start,
    gid_end,
    type_id_start,
    type_id_end,
    local_part_id,
    world_size,
    num_parts,
    cur_features,
    cur_global_ids,
):
    """This function is used to send/receive one feature for either nodes or
    edges of the input graph dataset.

    Parameters:
    -----------
    rank : int
        integer, unique id assigned to the current process
    data: dicitonary
        dictionry in which node or edge features are stored and this information
        is read from the appropriate node features file which belongs to the
        current process
    id_lookup : instance of DistLookupService
        instance of an implementation of dist. lookup service to retrieve values
        for keys
    feat_type : string
        this is used to distinguish which features are being exchanged. Please
        note that for nodes ownership is clearly defined and for edges it is
        always assumed that destination end point of the edge defines the
        ownership of that particular edge
    feat_key : string
        this string is used as a key in the dictionary to store features, as
        tensors, in local dictionaries
    featdata_key : numpy array
        features associated with this feature key being processed
    gid_start : int
        starting global_id, of either node or edge, for the feature data
    gid_end : int
        ending global_if, of either node or edge, for the feature data
    type_id_start : int
        starting type_id for the feature data
    type_id_end : int
        ending type_id for the feature data
    local_part_id : int
        integers used to the identify the local partition id used to locate
        data belonging to this partition
    world_size : int
        total number of processes created
    num_parts : int
        total number of partitions
    cur_features : dictionary
        dictionary to store the feature data which belongs to the current
        process
    cur_global_ids : dictionary
        dictionary to store global ids, of either nodes or edges, for which
        the features stored in the cur_features dictionary

    Returns:
    -------
    dictionary :
        a dictionary is returned where keys are type names and
        feature data are the values
    list :
        a dictionary of global_ids either nodes or edges whose features are
        received during the data shuffle process
    """
    # type_ids for this feature subset on the current rank
    gids_feat = np.arange(gid_start, gid_end)
    local_idx = np.arange(0, type_id_end - type_id_start)

    feats_per_rank = []
    global_id_per_rank = []

    tokens = feat_key.split("/")
    assert len(tokens) == 3
    local_feat_key = "/".join(tokens[:-1]) + "/" + str(local_part_id)

    logging.debug(
        f"[Rank: {rank} feature: {feat_key}, gid_start - {gid_start} and gid_end - {gid_end}"
    )

    # Get the partition ids for the range of global nids.
    if feat_type == constants.STR_NODE_FEATURES:
        # Retrieve the partition ids for the node features.
        # Each partition id will be in the range [0, num_parts).
        partid_slice = id_lookup.get_partition_ids(
            np.arange(gid_start, gid_end, dtype=np.int64)
        )
    else:
        # Edge data case.
        # Ownership is determined by the destination node.
        assert data is not None
        global_eids = np.arange(gid_start, gid_end, dtype=np.int64)
        if data[constants.GLOBAL_EID].shape[0] > 0:
            logging.debug(
                f"[Rank: {rank} disk read global eids - min - {np.amin(data[constants.GLOBAL_EID])}, max - {np.amax(data[constants.GLOBAL_EID])}, count - {data[constants.GLOBAL_EID].shape}"
            )

        # Now use `data` to extract destination nodes' global id
        # and use that to get the ownership
        common, idx1, idx2 = np.intersect1d(
            data[constants.GLOBAL_EID], global_eids, return_indices=True
        )
        assert (
            common.shape[0] == idx2.shape[0]
        ), f"Rank {rank}: {common.shape[0]} != {idx2.shape[0]}"
        assert (
            common.shape[0] == global_eids.shape[0]
        ), f"Rank {rank}: {common.shape[0]} != {global_eids.shape[0]}"

        global_dst_nids = data[constants.GLOBAL_DST_ID][idx1]
        assert np.all(global_eids == data[constants.GLOBAL_EID][idx1])
        partid_slice = id_lookup.get_partition_ids(global_dst_nids)

    # determine the shape of the feature-data
    # this is needed to so that ranks where feature-data is not present
    # should use the correct shape for sending the padded vector.
    # exchange length here.
    feat_dim_len = 0
    if featdata_key is not None:
        feat_dim_len = len(featdata_key.shape)
    all_lens = allgather_sizes(
        [feat_dim_len], world_size, num_parts, return_sizes=True
    )
    if all_lens[0] <= 0:
        logging.debug(
            f"[Rank: {rank} No process has any feature data to shuffle for {local_feat_key}"
        )
        return cur_features, cur_global_ids

    rank0_shape_len = all_lens[0]
    for idx in range(1, world_size):
        assert (all_lens[idx] == 0) or (all_lens[idx] == rank0_shape_len), (
            f"feature: {local_feat_key} shapes does not match "
            f"at rank - {idx} and rank - 0"
        )

    # exchange actual data here.
    if featdata_key is not None:
        logging.debug(f"Rank: {rank} {featdata_key.shape=}")
        feat_dims_dtype = list(featdata_key.shape)
        assert (
            len(featdata_key.shape) == 2 or len(featdata_key.shape) == 1
        ), f"We expect 1D or 2D tensors for features, got shape {featdata_key.shape}"
        # When a feature is 2-dim, the shape should match the feature dimension.
        if len(featdata_key.shape) == 2:
            feature_dimension = feat_dims_dtype[1]
        else:
            feature_dimension = 0
        feat_dims_dtype.append(DATA_TYPE_ID[featdata_key.dtype])
    else:
        feat_dims_dtype = list(np.zeros((rank0_shape_len), dtype=np.int64))
        feat_dims_dtype.append(DATA_TYPE_ID[torch.float32])
        feature_dimension = 0

    feature_dimension_tensor = torch.tensor([feature_dimension])
    dist.all_reduce(feature_dimension_tensor, op=dist.ReduceOp.MAX)
    feature_dimension = feature_dimension_tensor.item()

    logging.debug(f"Sending the feature shape information - {feat_dims_dtype}")
    all_dims_dtype = allgather_sizes(
        feat_dims_dtype, world_size, num_parts, return_sizes=True
    )

    for idx in range(world_size):
        cond = partid_slice == (idx + local_part_id * world_size)
        gids_per_partid = gids_feat[cond]
        local_idx_partid = local_idx[cond]

        if gids_per_partid.shape[0] == 0:
            assert len(all_dims_dtype) % world_size == 0
            dim_len = int(len(all_dims_dtype) / world_size)
            rank0_shape = list(np.zeros((dim_len - 1), dtype=np.int32))
            assert (
                len(rank0_shape) == 2 or len(rank0_shape) == 1
            ), f"We expect 1D or 2D tensors for features, got shape {rank0_shape}"
            # When a feature is 2-dim, the shape[1] (number of columns) should match the feature dimension.
            if len(rank0_shape) == 2:
                rank0_shape[1] = feature_dimension
            rank0_dtype = REV_DATA_TYPE_ID[
                all_dims_dtype[(dim_len - 1) : (dim_len)][0]
            ]
            data = torch.empty(rank0_shape, dtype=rank0_dtype)
            feats_per_rank.append(data)
            global_id_per_rank.append(torch.empty((0,), dtype=torch.int64))
        else:
            feats_per_rank.append(featdata_key[local_idx_partid])
            global_id_per_rank.append(
                torch.from_numpy(gids_per_partid).type(torch.int64)
            )
    for idx, tt in enumerate(feats_per_rank):
        logging.debug(
            f"[Rank: {rank} features shape - {tt.shape} and ids - {global_id_per_rank[idx].shape}"
        )

    # features (and global nids) per rank to be sent out are ready
    # for transmission, perform alltoallv here.
    output_feat_list = alltoallv_cpu(
        rank, world_size, feats_per_rank, retain_nones=False
    )
    output_id_list = alltoallv_cpu(
        rank, world_size, global_id_per_rank, retain_nones=False
    )
    logging.debug(
        f"[Rank : {rank} feats - {output_feat_list}, ids - {output_id_list}"
    )
    assert len(output_feat_list) == len(output_id_list), (
        "Length of feature list and id list are expected to be equal while "
        f"got {len(output_feat_list)} and {len(output_id_list)}."
    )

    # stitch node_features together to form one large feature tensor
    if len(output_feat_list) > 0:
        output_feat_list = torch.cat(output_feat_list)
        output_id_list = torch.cat(output_id_list)
        if local_feat_key in cur_features:
            temp = cur_features[local_feat_key]
            cur_features[local_feat_key] = torch.cat([temp, output_feat_list])
            temp = cur_global_ids[local_feat_key]
            cur_global_ids[local_feat_key] = torch.cat([temp, output_id_list])
        else:
            cur_features[local_feat_key] = output_feat_list
            cur_global_ids[local_feat_key] = output_id_list
    else:
        cur_features[local_feat_key] = torch.empty(
            (0, feature_dimension), dtype=torch.float32
        )
        cur_global_ids[local_feat_key] = torch.empty((0,), dtype=torch.int64)
    return cur_features, cur_global_ids
This function is used to send/receive one feature for either nodes or
    edges of the input graph dataset.

    Parameters:
    -----------
    rank : int
        integer, unique id assigned to the current process
    data: dicitonary
        dictionry in which node or edge features are stored and this information
        is read from the appropriate node features file which belongs to the
        current process
    id_lookup : instance of DistLookupService
        instance of an implementation of dist. lookup service to retrieve values
        for keys
    feat_type : string
        this is used to distinguish which features are being exchanged. Please
        note that for nodes ownership is clearly defined and for edges it is
        always assumed that destination end point of the edge defines the
        ownership of that particular edge
    feat_key : string
        this string is used as a key in the dictionary to store features, as
        tensors, in local dictionaries
    featdata_key : numpy array
        features associated with this feature key being processed
    gid_start : int
        starting global_id, of either node or edge, for the feature data
    gid_end : int
        ending global_if, of either node or edge, for the feature data
    type_id_start : int
        starting type_id for the feature data
    type_id_end : int
        ending type_id for the feature data
    local_part_id : int
        integers used to the identify the local partition id used to locate
        data belonging to this partition
    world_size : int
        total number of processes created
    num_parts : int
        total number of partitions
    cur_features : dictionary
        dictionary to store the feature data which belongs to the current
        process
    cur_global_ids : dictionary
        dictionary to store global ids, of either nodes or edges, for which
        the features stored in the cur_features dictionary

    Returns:
    -------
    dictionary :
        a dictionary is returned where keys are type names and
        feature data are the values
    list :
        a dictionary of global_ids either nodes or edges whose features are
        received during the data shuffle process
    
tools/distpartitioning/data_shuffle.py