def unique_and_compact(
    nodes: Union[
        List[torch.Tensor],
        Dict[str, List[torch.Tensor]],
    ],
    rank: int = 0,
    world_size: int = 1,
    async_op: bool = False,
):
    """
    Compact a list of nodes tensor. The `rank` and `world_size` parameters are
    relevant when using Cooperative Minibatching, which was initially proposed
    in `Deep Graph Library PR#4337<https://github.com/dmlc/dgl/pull/4337>`__ and
    was later first fully described in
    `Cooperative Minibatching in Graph Neural Networks
    <https://arxiv.org/abs/2310.12403>`__.
    Cooperation between the GPUs eliminates duplicate work performed across the
    GPUs due to the overlapping sampled k-hop neighborhoods of seed nodes when
    performing GNN minibatching.

    When `world_size` is greater than 1, then the given ids are partitioned
    between the available ranks. The ids corresponding to the given rank are
    guaranteed to come before the ids of other ranks. To do this, the
    partitioned ids are rotated backwards by the given rank so that the ids are
    ordered as: `[rank, rank + 1, world_size, 0, ..., rank - 1]`. This is
    supported only for Volta and later generation NVIDIA GPUs.

    Parameters
    ----------
    nodes : List[torch.Tensor] or Dict[str, List[torch.Tensor]]
        List of nodes for compacting.
        the unique_and_compact will be done per type
        - If `nodes` is a list of tensor: All the tensors will do unique and
        compact together, usually it is used for homogeneous graph.
        - If `nodes` is a list of dictionary: The keys should be node type and
        the values should be corresponding nodes, the unique and compact will
        be done per type, usually it is used for heterogeneous graph.
    rank : int
        The rank of the current process.
    world_size : int
        The number of processes.
    async_op: bool
        Boolean indicating whether the call is asynchronous. If so, the result
        can be obtained by calling wait on the returned future.

    Returns
    -------
    Tuple[unique_nodes, compacted_node_list, unique_nodes_offsets]
        The Unique nodes (per type) of all nodes in the input. And the compacted
        nodes list, where IDs inside are replaced with compacted node IDs.
        "Compacted node list" indicates that the node IDs in the input node
        list are replaced with mapped node IDs, where each type of node is
        mapped to a contiguous space of IDs ranging from 0 to N.
        The unique nodes offsets tensor partitions the unique_nodes tensor. Has
        size `world_size + 1` and `unique_nodes[offsets[i]: offsets[i + 1]]`
        belongs to the rank `(rank + i) % world_size`.
    """
    is_heterogeneous = isinstance(nodes, dict)

    if not is_heterogeneous:
        homo_ntype = "a"
        nodes = {homo_ntype: nodes}

    nums = {}
    concat_nodes, empties = [], []
    for ntype, nodes_of_type in nodes.items():
        nums[ntype] = [node.size(0) for node in nodes_of_type]
        concat_nodes.append(torch.cat(nodes_of_type))
        empties.append(concat_nodes[-1].new_empty(0))
    unique_fn = (
        torch.ops.graphbolt.unique_and_compact_batched_async
        if async_op
        else torch.ops.graphbolt.unique_and_compact_batched
    )
    results = unique_fn(concat_nodes, empties, empties, rank, world_size)

    class _Waiter:
        def __init__(self, future, ntypes, nums):
            self.future = future
            self.ntypes = ntypes
            self.nums = nums

        def wait(self):
            """Returns the stored value when invoked."""
            results = self.future.wait() if async_op else self.future
            ntypes = self.ntypes
            nums = self.nums
            # Ensure there is no memory leak.
            self.future = self.ntypes = self.nums = None

            unique, compacted, offsets = {}, {}, {}
            for ntype, result in zip(ntypes, results):
                (
                    unique[ntype],
                    concat_compacted,
                    _,
                    offsets[ntype],
                ) = result
                compacted[ntype] = list(concat_compacted.split(nums[ntype]))
            if is_heterogeneous:
                return unique, compacted, offsets
            else:
                return (
                    unique[homo_ntype],
                    compacted[homo_ntype],
                    offsets[homo_ntype],
                )

    post_processer = _Waiter(results, nodes.keys(), nums)
    if async_op:
        return post_processer
    else:
        return post_processer.wait()

    Compact a list of nodes tensor. The `rank` and `world_size` parameters are
    relevant when using Cooperative Minibatching, which was initially proposed
    in `Deep Graph Library PR#4337<https://github.com/dmlc/dgl/pull/4337>`__ and
    was later first fully described in
    `Cooperative Minibatching in Graph Neural Networks
    <https://arxiv.org/abs/2310.12403>`__.
    Cooperation between the GPUs eliminates duplicate work performed across the
    GPUs due to the overlapping sampled k-hop neighborhoods of seed nodes when
    performing GNN minibatching.

    When `world_size` is greater than 1, then the given ids are partitioned
    between the available ranks. The ids corresponding to the given rank are
    guaranteed to come before the ids of other ranks. To do this, the
    partitioned ids are rotated backwards by the given rank so that the ids are
    ordered as: `[rank, rank + 1, world_size, 0, ..., rank - 1]`. This is
    supported only for Volta and later generation NVIDIA GPUs.

    Parameters
    ----------
    nodes : List[torch.Tensor] or Dict[str, List[torch.Tensor]]
        List of nodes for compacting.
        the unique_and_compact will be done per type
        - If `nodes` is a list of tensor: All the tensors will do unique and
        compact together, usually it is used for homogeneous graph.
        - If `nodes` is a list of dictionary: The keys should be node type and
        the values should be corresponding nodes, the unique and compact will
        be done per type, usually it is used for heterogeneous graph.
    rank : int
        The rank of the current process.
    world_size : int
        The number of processes.
    async_op: bool
        Boolean indicating whether the call is asynchronous. If so, the result
        can be obtained by calling wait on the returned future.

    Returns
    -------
    Tuple[unique_nodes, compacted_node_list, unique_nodes_offsets]
        The Unique nodes (per type) of all nodes in the input. And the compacted
        nodes list, where IDs inside are replaced with compacted node IDs.
        "Compacted node list" indicates that the node IDs in the input node
        list are replaced with mapped node IDs, where each type of node is
        mapped to a contiguous space of IDs ranging from 0 to N.
        The unique nodes offsets tensor partitions the unique_nodes tensor. Has
        size `world_size + 1` and `unique_nodes[offsets[i]: offsets[i + 1]]`
        belongs to the rank `(rank + i) % world_size`.
    
python/dgl/graphbolt/internal/sample_utils.py