dgl.partition_graph_with_halo

dgl.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.

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.