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,83 @@
+from collections.abc import Mapping
+
+import dgl
+import numpy as np
+import pytest
+import torch
+
+
+def _create_homogeneous():
+    s = torch.randint(0, 200, (1000,))
+    d = torch.randint(0, 200, (1000,))
+    g = dgl.graph((s, d), num_nodes=200)
+    reverse_eids = torch.cat([torch.arange(1000, 2000), torch.arange(0, 1000)])
+    seed_edges = torch.arange(0, 1000)
+    return g, reverse_eids, seed_edges
+
+
+def _find_edges_to_exclude(g, pair_eids, degree_threshold):
+    src, dst = g.find_edges(pair_eids)
+    head_degree = g.in_degrees(src)
+    tail_degree = g.in_degrees(dst)
+    degree = torch.min(head_degree, tail_degree)
+    degree_mask = degree < degree_threshold
+    low_degree_pair_eids = pair_eids[degree_mask]
+    low_degree_pair_eids = torch.cat(
+        [low_degree_pair_eids, low_degree_pair_eids + 1000]
+    )
+    return low_degree_pair_eids
+
+
+@pytest.mark.parametrize("degree_threshold", [1, 2, 3, 4, 5])
+@pytest.mark.parametrize("batch_size", [1, 10, 50])
+def test_spot_target_excludes(degree_threshold, batch_size):
+    g, reverse_eids, seed_edges = _create_homogeneous()
+    sampler = dgl.dataloading.MultiLayerFullNeighborSampler(1)
+    low_degree_excluder = dgl.dataloading.SpotTarget(
+        g,
+        exclude="reverse_id",
+        degree_threshold=degree_threshold,
+        reverse_eids=reverse_eids,
+    )
+    sampler = dgl.dataloading.as_edge_prediction_sampler(
+        sampler,
+        exclude=low_degree_excluder,
+        negative_sampler=dgl.dataloading.negative_sampler.Uniform(1),
+    )
+    dataloader = dgl.dataloading.DataLoader(
+        g, seed_edges, sampler, batch_size=batch_size
+    )
+
+    for i, (input_nodes, pair_graph, neg_pair_graph, blocks) in enumerate(
+        dataloader
+    ):
+        if isinstance(blocks, list):
+            subg = blocks[0]
+        else:
+            subg = blocks
+        pair_eids = pair_graph.edata[dgl.EID]
+        block_eids = subg.edata[dgl.EID]
+        edges_to_exclude = _find_edges_to_exclude(
+            g, pair_eids, degree_threshold
+        )
+        if edges_to_exclude is None:
+            continue
+        edges_to_exclude = dgl.utils.recursive_apply(
+            edges_to_exclude, lambda x: x.cpu().numpy()
+        )
+        block_eids = dgl.utils.recursive_apply(
+            block_eids, lambda x: x.cpu().numpy()
+        )
+
+        if isinstance(edges_to_exclude, Mapping):
+            for k in edges_to_exclude.keys():
+                assert not np.isin(edges_to_exclude[k], block_eids[k]).any()
+        else:
+            assert not np.isin(edges_to_exclude, block_eids).any()
+
+        if i == 10:
+            break
+
+
+if __name__ == "__main__":
+    test_spot_target_excludes(degree_threshold=2, batch_size=10)