chore: import upstream snapshot with attribution

2026-07-13 13:35:51 +08:00
commit c36a561cd8
2172 changed files with 455595 additions and 0 deletions
@@ -0,0 +1,64 @@
+"""Edge coarsening procedure used in Metis and Graclus, for pytorch"""
+# pylint: disable=no-member, invalid-name, W0613
+from .. import remove_self_loop
+from .capi import _neighbor_matching
+
+__all__ = ["neighbor_matching"]
+
+
+def neighbor_matching(graph, e_weights=None, relabel_idx=True):
+    r"""
+    Description
+    -----------
+    The neighbor matching procedure of edge coarsening in
+    `Metis <http://cacs.usc.edu/education/cs653/Karypis-METIS-SIAMJSC98.pdf>`__
+    and
+    `Graclus <https://www.cs.utexas.edu/users/inderjit/public_papers/multilevel_pami.pdf>`__
+    for homogeneous graph coarsening. This procedure keeps picking an unmarked
+    vertex and matching it with one its unmarked neighbors (that maximizes its
+    edge weight) until no match can be done.
+
+    If no edge weight is given, this procedure will randomly pick neighbor for each
+    vertex.
+
+    The GPU implementation is based on `A GPU Algorithm for Greedy Graph Matching
+    <http://www.staff.science.uu.nl/~bisse101/Articles/match12.pdf>`__
+
+    NOTE: The input graph must be bi-directed (undirected) graph. Call :obj:`dgl.to_bidirected`
+          if you are not sure your graph is bi-directed.
+
+    Parameters
+    ----------
+    graph : DGLGraph
+        The input homogeneous graph.
+    edge_weight : torch.Tensor, optional
+        The edge weight tensor holding non-negative scalar weight for each edge.
+        default: :obj:`None`
+    relabel_idx : bool, optional
+        If true, relabel resulting node labels to have consecutive node ids.
+        default: :obj:`True`
+
+    Examples
+    --------
+    The following example uses PyTorch backend.
+
+    >>> import torch, dgl
+    >>> from dgl.geometry import neighbor_matching
+    >>>
+    >>> g = dgl.graph(([0, 1, 1, 2], [1, 0, 2, 1]))
+    >>> res = neighbor_matching(g)
+        tensor([0, 1, 1])
+    """
+    assert (
+        graph.is_homogeneous
+    ), "The graph used in graph node matching must be homogeneous"
+    if e_weights is not None:
+        graph.edata["e_weights"] = e_weights
+        graph = remove_self_loop(graph)
+        e_weights = graph.edata["e_weights"]
+        graph.edata.pop("e_weights")
+    else:
+        graph = remove_self_loop(graph)
+    return _neighbor_matching(
+        graph._graph, graph.num_nodes(), e_weights, relabel_idx
+    )