torch_geometric.nn.pool.graclus

graclus(edge_index: Tensor, weight: Optional[Tensor] = None, num_nodes: Optional[int] = None)[source]

A greedy clustering algorithm from the “Weighted Graph Cuts without Eigenvectors: A Multilevel Approach” paper of picking an unmarked vertex and matching it with one of its unmarked neighbors (that maximizes its edge weight). The GPU algorithm is adapted from the “A GPU Algorithm for Greedy Graph Matching” paper.

Parameters:

  • edge_index (torch.Tensor) – The edge indices.

  • weight (torch.Tensor, optional) – One-dimensional edge weights. (default: None)

  • num_nodes (int, optional) – The number of nodes, i.e. max_val + 1 of edge_index. (default: None)

Return type:

torch.Tensor