partition_graph_with_halo(g, node_part, extra_cached_hops, reshuffle=False)[source]¶
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.
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.
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.