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,123 @@
+"""Torch Module for APPNPConv"""
+# pylint: disable= no-member, arguments-differ, invalid-name
+import torch as th
+from torch import nn
+
+from .... import function as fn
+from .graphconv import EdgeWeightNorm
+
+
+class APPNPConv(nn.Module):
+    r"""Approximate Personalized Propagation of Neural Predictions layer from `Predict then
+    Propagate: Graph Neural Networks meet Personalized PageRank
+    <https://arxiv.org/pdf/1810.05997.pdf>`__
+
+    .. math::
+        H^{0} &= X
+
+        H^{l+1} &= (1-\alpha)\left(\tilde{D}^{-1/2}
+        \tilde{A} \tilde{D}^{-1/2} H^{l}\right) + \alpha H^{0}
+
+    where :math:`\tilde{A}` is :math:`A` + :math:`I`.
+
+    Parameters
+    ----------
+    k : int
+        The number of iterations :math:`K`.
+    alpha : float
+        The teleport probability :math:`\alpha`.
+    edge_drop : float, optional
+        The dropout rate on edges that controls the
+        messages received by each node. Default: ``0``.
+
+    Example
+    -------
+    >>> import dgl
+    >>> import numpy as np
+    >>> import torch as th
+    >>> from dgl.nn import APPNPConv
+    >>>
+    >>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
+    >>> feat = th.ones(6, 10)
+    >>> conv = APPNPConv(k=3, alpha=0.5)
+    >>> res = conv(g, feat)
+    >>> print(res)
+    tensor([[0.8536, 0.8536, 0.8536, 0.8536, 0.8536, 0.8536, 0.8536, 0.8536, 0.8536,
+            0.8536],
+            [0.9268, 0.9268, 0.9268, 0.9268, 0.9268, 0.9268, 0.9268, 0.9268, 0.9268,
+            0.9268],
+            [0.9634, 0.9634, 0.9634, 0.9634, 0.9634, 0.9634, 0.9634, 0.9634, 0.9634,
+            0.9634],
+            [0.9268, 0.9268, 0.9268, 0.9268, 0.9268, 0.9268, 0.9268, 0.9268, 0.9268,
+            0.9268],
+            [0.9634, 0.9634, 0.9634, 0.9634, 0.9634, 0.9634, 0.9634, 0.9634, 0.9634,
+            0.9634],
+            [0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000,
+            0.5000]])
+    """
+
+    def __init__(self, k, alpha, edge_drop=0.0):
+        super(APPNPConv, self).__init__()
+        self._k = k
+        self._alpha = alpha
+        self.edge_drop = nn.Dropout(edge_drop)
+
+    def forward(self, graph, feat, edge_weight=None):
+        r"""
+
+        Description
+        -----------
+        Compute APPNP layer.
+
+        Parameters
+        ----------
+        graph : DGLGraph
+            The graph.
+        feat : torch.Tensor
+            The input feature of shape :math:`(N, *)`. :math:`N` is the
+            number of nodes, and :math:`*` could be of any shape.
+        edge_weight: torch.Tensor, optional
+            edge_weight to use in the message passing process. This is equivalent to
+            using weighted adjacency matrix in the equation above, and
+            :math:`\tilde{D}^{-1/2}\tilde{A} \tilde{D}^{-1/2}`
+            is based on :class:`dgl.nn.pytorch.conv.graphconv.EdgeWeightNorm`.
+
+        Returns
+        -------
+        torch.Tensor
+            The output feature of shape :math:`(N, *)` where :math:`*`
+            should be the same as input shape.
+        """
+        with graph.local_scope():
+            if edge_weight is None:
+                src_norm = th.pow(
+                    graph.out_degrees().to(feat).clamp(min=1), -0.5
+                )
+                shp = src_norm.shape + (1,) * (feat.dim() - 1)
+                src_norm = th.reshape(src_norm, shp).to(feat.device)
+                dst_norm = th.pow(
+                    graph.in_degrees().to(feat).clamp(min=1), -0.5
+                )
+                shp = dst_norm.shape + (1,) * (feat.dim() - 1)
+                dst_norm = th.reshape(dst_norm, shp).to(feat.device)
+            else:
+                edge_weight = EdgeWeightNorm("both")(graph, edge_weight)
+            feat_0 = feat
+            for _ in range(self._k):
+                # normalization by src node
+                if edge_weight is None:
+                    feat = feat * src_norm
+                graph.ndata["h"] = feat
+                w = (
+                    th.ones(graph.num_edges(), 1)
+                    if edge_weight is None
+                    else edge_weight
+                )
+                graph.edata["w"] = self.edge_drop(w).to(feat.device)
+                graph.update_all(fn.u_mul_e("h", "w", "m"), fn.sum("m", "h"))
+                feat = graph.ndata.pop("h")
+                # normalization by dst node
+                if edge_weight is None:
+                    feat = feat * dst_norm
+                feat = (1 - self._alpha) * feat + self._alpha * feat_0
+            return feat