def partition_graph_with_halo(g, node_part, extra_cached_hops, reshuffle=False):
    """Partition a graph.

    Based on the given node assignments for each partition, the function splits
    the input graph into subgraphs. A subgraph may contain HALO nodes which does
    not belong to the partition of a subgraph but are connected to the nodes
    in the partition within a fixed number of hops.

    If `reshuffle` is turned on, the function reshuffles node IDs and edge IDs
    of the input graph before partitioning. After reshuffling, all nodes and edges
    in a partition fall in a contiguous ID range in the input graph.
    The partitioend subgraphs have node data 'orig_id', which stores the node IDs
    in the original input graph.

    Parameters
    ------------
    g: DGLGraph
        The graph to be partitioned
    node_part: 1D tensor
        Specify which partition a node is assigned to. The length of this tensor
        needs to be the same as the number of nodes of the graph. Each element
        indicates the partition ID of a node.
    extra_cached_hops: int
        The number of hops a HALO node can be accessed.
    reshuffle : bool
        Resuffle nodes so that nodes in the same partition are in the same ID range.

    Returns
    --------
    a dict of DGLGraphs
        The key is the partition ID and the value is the DGLGraph of the partition.
    Tensor
        1D tensor that stores the mapping between the reshuffled node IDs and
        the original node IDs if 'reshuffle=True'. Otherwise, return None.
    Tensor
        1D tensor that stores the mapping between the reshuffled edge IDs and
        the original edge IDs if 'reshuffle=True'. Otherwise, return None.
    """
    assert len(node_part) == g.num_nodes()
    if reshuffle:
        g, node_part = reshuffle_graph(g, node_part)
        orig_nids = g.ndata["orig_id"]
        orig_eids = g.edata["orig_id"]

    node_part = utils.toindex(node_part)
    start = time.time()
    subgs = _CAPI_DGLPartitionWithHalo_Hetero(
        g._graph, node_part.todgltensor(), extra_cached_hops
    )
    # g is no longer needed. Free memory.
    g = None
    print("Split the graph: {:.3f} seconds".format(time.time() - start))
    subg_dict = {}
    node_part = node_part.tousertensor()
    start = time.time()

    # This function determines whether an edge belongs to a partition.
    # An edge is assigned to a partition based on its destination node. If its destination node
    # is assigned to a partition, we assign the edge to the partition as well.
    def get_inner_edge(subg, inner_node):
        inner_edge = F.zeros((subg.num_edges(),), F.int8, F.cpu())
        inner_nids = F.nonzero_1d(inner_node)
        # TODO(zhengda) we need to fix utils.toindex() to avoid the dtype cast below.
        inner_nids = F.astype(inner_nids, F.int64)
        inner_eids = subg.in_edges(inner_nids, form="eid")
        inner_edge = F.scatter_row(
            inner_edge,
            inner_eids,
            F.ones((len(inner_eids),), F.dtype(inner_edge), F.cpu()),
        )
        return inner_edge

    # This creaets a subgraph from subgraphs returned from the CAPI above.
    def create_subgraph(subg, induced_nodes, induced_edges, inner_node):
        subg1 = DGLGraph(gidx=subg.graph, ntypes=["_N"], etypes=["_E"])
        # If IDs are shuffled, we should shuffled edges. This will help us collect edge data
        # from the distributed graph after training.
        if reshuffle:
            # When we shuffle edges, we need to make sure that the inner edges are assigned with
            # contiguous edge IDs and their ID range starts with 0. In other words, we want to
            # place these edge IDs in the front of the edge list. To ensure that, we add the IDs
            # of outer edges with a large value, so we will get the sorted list as we want.
            max_eid = F.max(induced_edges[0], 0) + 1
            inner_edge = get_inner_edge(subg1, inner_node)
            eid = F.astype(induced_edges[0], F.int64) + max_eid * F.astype(
                inner_edge == 0, F.int64
            )

            _, index = F.sort_1d(eid)
            subg1 = edge_subgraph(subg1, index, relabel_nodes=False)
            subg1.ndata[NID] = induced_nodes[0]
            subg1.edata[EID] = F.gather_row(induced_edges[0], index)
        else:
            subg1.ndata[NID] = induced_nodes[0]
            subg1.edata[EID] = induced_edges[0]
        return subg1

    for i, subg in enumerate(subgs):
        inner_node = _get_halo_heterosubgraph_inner_node(subg)
        inner_node = F.zerocopy_from_dlpack(inner_node.to_dlpack())
        subg = create_subgraph(
            subg, subg.induced_nodes, subg.induced_edges, inner_node
        )
        subg.ndata["inner_node"] = inner_node
        subg.ndata["part_id"] = F.gather_row(node_part, subg.ndata[NID])
        if reshuffle:
            subg.ndata["orig_id"] = F.gather_row(orig_nids, subg.ndata[NID])
            subg.edata["orig_id"] = F.gather_row(orig_eids, subg.edata[EID])

        if extra_cached_hops >= 1:
            inner_edge = get_inner_edge(subg, inner_node)
        else:
            inner_edge = F.ones((subg.num_edges(),), F.int8, F.cpu())
        subg.edata["inner_edge"] = inner_edge
        subg_dict[i] = subg
    print("Construct subgraphs: {:.3f} seconds".format(time.time() - start))
    if reshuffle:
        return subg_dict, orig_nids, orig_eids
    else:
        return subg_dict, None, None
Partition a graph.

    Based on the given node assignments for each partition, the function splits
    the input graph into subgraphs. A subgraph may contain HALO nodes which does
    not belong to the partition of a subgraph but are connected to the nodes
    in the partition within a fixed number of hops.

    If `reshuffle` is turned on, the function reshuffles node IDs and edge IDs
    of the input graph before partitioning. After reshuffling, all nodes and edges
    in a partition fall in a contiguous ID range in the input graph.
    The partitioend subgraphs have node data 'orig_id', which stores the node IDs
    in the original input graph.

    Parameters
    ------------
    g: DGLGraph
        The graph to be partitioned
    node_part: 1D tensor
        Specify which partition a node is assigned to. The length of this tensor
        needs to be the same as the number of nodes of the graph. Each element
        indicates the partition ID of a node.
    extra_cached_hops: int
        The number of hops a HALO node can be accessed.
    reshuffle : bool
        Resuffle nodes so that nodes in the same partition are in the same ID range.

    Returns
    --------
    a dict of DGLGraphs
        The key is the partition ID and the value is the DGLGraph of the partition.
    Tensor
        1D tensor that stores the mapping between the reshuffled node IDs and
        the original node IDs if 'reshuffle=True'. Otherwise, return None.
    Tensor
        1D tensor that stores the mapping between the reshuffled edge IDs and
        the original edge IDs if 'reshuffle=True'. Otherwise, return None.
    
python/dgl/partition.py