pyg_lib.partition

metis(rowptr: Tensor, col: Tensor, num_partitions: int, node_weight: Optional[Tensor] = None, edge_weight: Optional[Tensor] = None, recursive: bool = False) → Tensor[source]

Clusters/partitions a graph into multiple partitions via METIS, as motivated by the “Cluster-GCN: An Efficient Algorithm for Training Deep and Large Graph Convolutional Networks” paper.

Parameters:

  • rowptr (Tensor) – Compressed source node indices.

  • col (Tensor) – Target node indices.

  • num_partitions (int) – The number of partitions.

  • node_weight (Optional[Tensor], default: None) – The node weights.

  • edge_weight (Optional[Tensor], default: None) – The edge weights.

  • recursive (bool, default: False) – If set to True, will use multilevel recursive bisection instead of multilevel k-way partitioning.

Returns:

Tensor – A vector that assings each node to a partition.