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,105 @@
+"""Sampling utilities"""
+from collections.abc import Mapping
+
+import numpy as np
+
+from .. import backend as F, transforms, utils
+from ..base import EID
+
+from ..utils import recursive_apply, recursive_apply_pair
+
+
+def _locate_eids_to_exclude(frontier_parent_eids, exclude_eids):
+    """Find the edges whose IDs in parent graph appeared in exclude_eids.
+
+    Note that both arguments are numpy arrays or numpy dicts.
+    """
+    if not isinstance(frontier_parent_eids, Mapping):
+        return np.isin(frontier_parent_eids, exclude_eids).nonzero()[0]
+    result = {}
+    for k, v in frontier_parent_eids.items():
+        if k in exclude_eids:
+            result[k] = np.isin(v, exclude_eids[k]).nonzero()[0]
+    return recursive_apply(result, F.zerocopy_from_numpy)
+
+
+class EidExcluder(object):
+    """Class that finds the edges whose IDs in parent graph appeared in exclude_eids.
+
+    The edge IDs can be both CPU and GPU tensors.
+    """
+
+    def __init__(self, exclude_eids):
+        device = None
+        if isinstance(exclude_eids, Mapping):
+            for _, v in exclude_eids.items():
+                if device is None:
+                    device = F.context(v)
+                    break
+        else:
+            device = F.context(exclude_eids)
+        self._exclude_eids = None
+        self._filter = None
+
+        if device == F.cpu():
+            # TODO(nv-dlasalle): Once Filter is implemented for the CPU, we
+            # should just use that irregardless of the device.
+            self._exclude_eids = (
+                recursive_apply(exclude_eids, F.zerocopy_to_numpy)
+                if exclude_eids is not None
+                else None
+            )
+        else:
+            self._filter = recursive_apply(exclude_eids, utils.Filter)
+
+    def _find_indices(self, parent_eids):
+        """Find the set of edge indices to remove."""
+        if self._exclude_eids is not None:
+            parent_eids_np = recursive_apply(parent_eids, F.zerocopy_to_numpy)
+            return _locate_eids_to_exclude(parent_eids_np, self._exclude_eids)
+        else:
+            assert self._filter is not None
+            func = lambda x, y: x.find_included_indices(y)
+            return recursive_apply_pair(self._filter, parent_eids, func)
+
+    def __call__(self, frontier, weights=None):
+        parent_eids = frontier.edata[EID]
+        located_eids = self._find_indices(parent_eids)
+
+        if not isinstance(located_eids, Mapping):
+            # (BarclayII) If frontier already has a EID field and located_eids is empty,
+            # the returned graph will keep EID intact.  Otherwise, EID will change
+            # to the mapping from the new graph to the old frontier.
+            # So we need to test if located_eids is empty, and do the remapping ourselves.
+            if len(located_eids) > 0:
+                frontier = transforms.remove_edges(
+                    frontier, located_eids, store_ids=True
+                )
+                if (
+                    weights is not None
+                    and weights[0].shape[0] == frontier.num_edges()
+                ):
+                    weights[0] = F.gather_row(weights[0], frontier.edata[EID])
+                frontier.edata[EID] = F.gather_row(
+                    parent_eids, frontier.edata[EID]
+                )
+        else:
+            # (BarclayII) remove_edges only accepts removing one type of edges,
+            # so I need to keep track of the edge IDs left one by one.
+            new_eids = parent_eids.copy()
+            for i, (k, v) in enumerate(located_eids.items()):
+                if len(v) > 0:
+                    frontier = transforms.remove_edges(
+                        frontier, v, etype=k, store_ids=True
+                    )
+                    new_eids[k] = F.gather_row(
+                        parent_eids[k], frontier.edges[k].data[EID]
+                    )
+                    if weights is not None and weights[i].shape[
+                        0
+                    ] == frontier.num_edges(k):
+                        weights[i] = F.gather_row(
+                            weights[i], frontier.edges[k].data[EID]
+                        )
+            frontier.edata[EID] = new_eids
+        return frontier if weights is None else (frontier, weights)