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 (torch.Tensor) – Compressed source node indices.

  • col (torch.Tensor) – Target node indices.

  • num_partitions (int) – The number of partitions.

  • node_weight (torch.Tensor, optional) – Optional node weights. (default: None)

  • edge_weight (torch.Tensor, optional) – Optional edge weights. (default: None)

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

Returns:

A vector that assings each node to a partition.

Return type:

torch.Tensor