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,57 @@
+"""An abstract base class for cugraph-ops nn module."""
+import torch
+from torch import nn
+
+
+class CuGraphBaseConv(nn.Module):
+    r"""An abstract base class for cugraph-ops nn module."""
+
+    def __init__(self):
+        super().__init__()
+        self._cached_offsets_fg = None
+
+    def reset_parameters(self):
+        r"""Resets all learnable parameters of the module."""
+        raise NotImplementedError
+
+    def forward(self, *args):
+        r"""Runs the forward pass of the module."""
+        raise NotImplementedError
+
+    def pad_offsets(self, offsets: torch.Tensor, size: int) -> torch.Tensor:
+        r"""Pad zero-in-degree nodes to the end of offsets to reach size.
+
+        cugraph-ops often provides two variants of aggregation functions for a
+        specific model: one intended for sampled-graph use cases, one for
+        full-graph ones. The former is in general more performant, however, it
+        only works when the sample size (the max of in-degrees) is small (<200),
+        due to the limit of GPU shared memory. For graphs with a larger max
+        in-degree, we need to fall back to the full-graph option, which requires
+        to convert a DGL block to a full graph. With the csc-representation,
+        this is equivalent to pad zero-in-degree nodes to the end of the offsets
+        array (also called indptr or colptr).
+
+        Parameters
+        ----------
+        offsets :
+            The (monotonically increasing) index pointer array in a CSC-format
+            graph.
+        size : int
+            The length of offsets after padding.
+
+        Returns
+        -------
+        torch.Tensor
+            The augmented offsets array.
+        """
+        if self._cached_offsets_fg is None:
+            self._cached_offsets_fg = torch.empty(
+                size, dtype=offsets.dtype, device=offsets.device
+            )
+        elif self._cached_offsets_fg.numel() < size:
+            self._cached_offsets_fg.resize_(size)
+
+        self._cached_offsets_fg[: offsets.numel()] = offsets
+        self._cached_offsets_fg[offsets.numel() : size] = offsets[-1]
+
+        return self._cached_offsets_fg[:size]