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,7 @@
+import backend as F
+import pytest
+
+parametrize_idtype = pytest.mark.parametrize("idtype", [F.int32, F.int64])
+
+from .checks import *
+from .graph_cases import get_cases
@@ -0,0 +1,124 @@
+import backend as F
+
+import dgl
+import pytest
+from dgl.base import is_internal_column
+
+__all__ = [
+    "check_fail",
+    "assert_is_identical",
+    "assert_is_identical_hetero",
+    "check_graph_equal",
+]
+
+
+def check_fail(fn, *args, **kwargs):
+    try:
+        fn(*args, **kwargs)
+        return False
+    except:
+        return True
+
+
+def assert_is_identical(g, g2):
+    assert g.num_nodes() == g2.num_nodes()
+    src, dst = g.all_edges(order="eid")
+    src2, dst2 = g2.all_edges(order="eid")
+    assert F.array_equal(src, src2)
+    assert F.array_equal(dst, dst2)
+
+    assert len(g.ndata) == len(g2.ndata)
+    assert len(g.edata) == len(g2.edata)
+    for k in g.ndata:
+        assert F.allclose(g.ndata[k], g2.ndata[k])
+    for k in g.edata:
+        assert F.allclose(g.edata[k], g2.edata[k])
+
+
+def assert_is_identical_hetero(g, g2, ignore_internal_data=False):
+    assert g.ntypes == g2.ntypes
+    assert g.canonical_etypes == g2.canonical_etypes
+
+    # check if two metagraphs are identical
+    for edges, features in g.metagraph().edges(keys=True).items():
+        assert g2.metagraph().edges(keys=True)[edges] == features
+
+    # check if node ID spaces and feature spaces are equal
+    for ntype in g.ntypes:
+        assert g.num_nodes(ntype) == g2.num_nodes(ntype)
+        if ignore_internal_data:
+            for k in list(g.nodes[ntype].data.keys()):
+                if is_internal_column(k):
+                    del g.nodes[ntype].data[k]
+            for k in list(g2.nodes[ntype].data.keys()):
+                if is_internal_column(k):
+                    del g2.nodes[ntype].data[k]
+        assert len(g.nodes[ntype].data) == len(g2.nodes[ntype].data)
+        for k in g.nodes[ntype].data:
+            assert F.allclose(g.nodes[ntype].data[k], g2.nodes[ntype].data[k])
+
+    # check if edge ID spaces and feature spaces are equal
+    for etype in g.canonical_etypes:
+        src, dst = g.all_edges(etype=etype, order="eid")
+        src2, dst2 = g2.all_edges(etype=etype, order="eid")
+        assert F.array_equal(src, src2)
+        assert F.array_equal(dst, dst2)
+        if ignore_internal_data:
+            for k in list(g.edges[etype].data.keys()):
+                if is_internal_column(k):
+                    del g.edges[etype].data[k]
+            for k in list(g2.edges[etype].data.keys()):
+                if is_internal_column(k):
+                    del g2.edges[etype].data[k]
+        assert len(g.edges[etype].data) == len(g2.edges[etype].data)
+        for k in g.edges[etype].data:
+            assert F.allclose(g.edges[etype].data[k], g2.edges[etype].data[k])
+
+
+def check_graph_equal(g1, g2, *, check_idtype=True, check_feature=True):
+    assert g1.device == g2.device
+    if check_idtype:
+        assert g1.idtype == g2.idtype
+    assert g1.ntypes == g2.ntypes
+    assert g1.etypes == g2.etypes
+    assert g1.srctypes == g2.srctypes
+    assert g1.dsttypes == g2.dsttypes
+    assert g1.canonical_etypes == g2.canonical_etypes
+    assert g1.batch_size == g2.batch_size
+
+    # check if two metagraphs are identical
+    for edges, features in g1.metagraph().edges(keys=True).items():
+        assert g2.metagraph().edges(keys=True)[edges] == features
+
+    for nty in g1.ntypes:
+        assert g1.num_nodes(nty) == g2.num_nodes(nty)
+        assert F.allclose(g1.batch_num_nodes(nty), g2.batch_num_nodes(nty))
+    for ety in g1.canonical_etypes:
+        assert g1.num_edges(ety) == g2.num_edges(ety)
+        assert F.allclose(g1.batch_num_edges(ety), g2.batch_num_edges(ety))
+        src1, dst1, eid1 = g1.edges(etype=ety, form="all")
+        src2, dst2, eid2 = g2.edges(etype=ety, form="all")
+        if check_idtype:
+            assert F.allclose(src1, src2)
+            assert F.allclose(dst1, dst2)
+            assert F.allclose(eid1, eid2)
+        else:
+            assert F.allclose(src1, F.astype(src2, g1.idtype))
+            assert F.allclose(dst1, F.astype(dst2, g1.idtype))
+            assert F.allclose(eid1, F.astype(eid2, g1.idtype))
+
+    if check_feature:
+        for nty in g1.ntypes:
+            if g1.num_nodes(nty) == 0:
+                continue
+            for feat_name in g1.nodes[nty].data.keys():
+                assert F.allclose(
+                    g1.nodes[nty].data[feat_name], g2.nodes[nty].data[feat_name]
+                )
+        for ety in g1.canonical_etypes:
+            if g1.num_edges(ety) == 0:
+                continue
+            for feat_name in g2.edges[ety].data.keys():
+                assert F.allclose(
+                    g1.edges[ety].data[feat_name], g2.edges[ety].data[feat_name]
+                )
@@ -0,0 +1,266 @@
+from collections import defaultdict
+
+import backend as F
+import dgl
+import networkx as nx
+import numpy as np
+import scipy.sparse as ssp
+
+case_registry = defaultdict(list)
+
+
+def register_case(labels):
+    def wrapper(fn):
+        for lbl in labels:
+            case_registry[lbl].append(fn)
+        fn.__labels__ = labels
+        return fn
+
+    return wrapper
+
+
+def get_cases(labels=None, exclude=[]):
+    """Get all graph instances of the given labels."""
+    cases = set()
+    if labels is None:
+        # get all the cases
+        labels = case_registry.keys()
+    for lbl in labels:
+        for case in case_registry[lbl]:
+            if not any([l in exclude for l in case.__labels__]):
+                cases.add(case)
+    return [fn() for fn in cases]
+
+
+@register_case(["bipartite", "zero-degree"])
+def bipartite1():
+    return dgl.heterograph(
+        {("_U", "_E", "_V"): ([0, 0, 0, 2, 2, 3], [0, 1, 4, 1, 4, 3])}
+    )
+
+
+@register_case(["bipartite"])
+def bipartite_full():
+    return dgl.heterograph(
+        {
+            ("_U", "_E", "_V"): (
+                [0, 0, 0, 0, 1, 1, 1, 1],
+                [0, 1, 2, 3, 0, 1, 2, 3],
+            )
+        }
+    )
+
+
+@register_case(["homo"])
+def graph0():
+    return dgl.graph(
+        (
+            [0, 0, 0, 1, 1, 2, 2, 2, 3, 3, 3, 4, 6, 6, 7, 8, 9],
+            [4, 5, 1, 2, 4, 7, 9, 8, 6, 4, 1, 0, 1, 0, 2, 3, 5],
+        )
+    )
+
+
+@register_case(["homo", "zero-degree", "homo-zero-degree"])
+def bipartite1():
+    return dgl.graph(([0, 0, 0, 2, 2, 3], [0, 1, 4, 1, 4, 3]))
+
+
+@register_case(["homo", "has_feature"])
+def graph1():
+    g = dgl.graph(
+        (
+            [0, 0, 0, 1, 1, 2, 2, 2, 3, 3, 3, 4, 6, 6, 7, 8, 9],
+            [4, 5, 1, 2, 4, 7, 9, 8, 6, 4, 1, 0, 1, 0, 2, 3, 5],
+        ),
+        device=F.cpu(),
+    )
+    g.ndata["h"] = F.copy_to(F.randn((g.num_nodes(), 2)), F.cpu())
+    g.edata["w"] = F.copy_to(F.randn((g.num_edges(), 3)), F.cpu())
+    return g
+
+
+@register_case(["homo", "has_scalar_e_feature"])
+def graph1():
+    g = dgl.graph(
+        (
+            [0, 0, 0, 1, 1, 2, 2, 2, 3, 3, 3, 4, 6, 6, 7, 8, 9],
+            [4, 5, 1, 2, 4, 7, 9, 8, 6, 4, 1, 0, 1, 0, 2, 3, 5],
+        ),
+        device=F.cpu(),
+    )
+    g.ndata["h"] = F.copy_to(F.randn((g.num_nodes(), 2)), F.cpu())
+    g.edata["scalar_w"] = F.copy_to(F.abs(F.randn((g.num_edges(),))), F.cpu())
+    return g
+
+
+@register_case(["homo", "row_sorted"])
+def graph2():
+    return dgl.graph(
+        (
+            [0, 0, 0, 1, 1, 2, 2, 2, 3, 3, 3, 4, 6, 6, 7, 8, 9],
+            [4, 5, 1, 2, 4, 7, 9, 8, 6, 4, 1, 0, 1, 0, 2, 3, 5],
+        ),
+        row_sorted=True,
+    )
+
+
+@register_case(["homo", "row_sorted", "col_sorted"])
+def graph3():
+    return dgl.graph(
+        (
+            [0, 0, 0, 1, 1, 2, 2, 2, 3, 3, 3, 4, 6, 6, 7, 8, 9],
+            [1, 4, 5, 2, 4, 7, 8, 9, 1, 4, 6, 0, 0, 1, 2, 3, 5],
+        ),
+        row_sorted=True,
+        col_sorted=True,
+    )
+
+
+@register_case(["hetero", "has_feature"])
+def heterograph0():
+    g = dgl.heterograph(
+        {
+            ("user", "plays", "game"): ([0, 1, 1, 2], [0, 0, 1, 1]),
+            ("developer", "develops", "game"): ([0, 1], [0, 1]),
+        },
+        device=F.cpu(),
+    )
+    g.nodes["user"].data["h"] = F.copy_to(
+        F.randn((g.num_nodes("user"), 3)), F.cpu()
+    )
+    g.nodes["game"].data["h"] = F.copy_to(
+        F.randn((g.num_nodes("game"), 2)), F.cpu()
+    )
+    g.nodes["developer"].data["h"] = F.copy_to(
+        F.randn((g.num_nodes("developer"), 3)), F.cpu()
+    )
+    g.edges["plays"].data["h"] = F.copy_to(
+        F.randn((g.num_edges("plays"), 1)), F.cpu()
+    )
+    g.edges["develops"].data["h"] = F.copy_to(
+        F.randn((g.num_edges("develops"), 5)), F.cpu()
+    )
+    return g
+
+
+@register_case(["batched", "homo"])
+def batched_graph0():
+    g1 = dgl.add_self_loop(dgl.graph(([0, 1, 2], [1, 2, 3])))
+    g2 = dgl.add_self_loop(dgl.graph(([1, 1], [2, 0])))
+    g3 = dgl.add_self_loop(dgl.graph(([0], [1])))
+    return dgl.batch([g1, g2, g3])
+
+
+@register_case(["block", "bipartite", "block-bipartite"])
+def block_graph0():
+    g = dgl.graph(([2, 3, 4], [5, 6, 7]), num_nodes=100)
+    g = g.to(F.cpu())
+    return dgl.to_block(g)
+
+
+@register_case(["block"])
+def block_graph1():
+    g = dgl.heterograph(
+        {
+            ("user", "plays", "game"): ([0, 1, 2], [1, 1, 0]),
+            ("user", "likes", "game"): ([1, 2, 3], [0, 0, 2]),
+            ("store", "sells", "game"): ([0, 1, 1], [0, 1, 2]),
+        },
+        device=F.cpu(),
+    )
+    return dgl.to_block(g)
+
+
+@register_case(["clique"])
+def clique():
+    g = dgl.graph(([0, 0, 0, 1, 1, 1, 2, 2, 2], [0, 1, 2, 0, 1, 2, 0, 1, 2]))
+    return g
+
+
+def random_dglgraph(size):
+    return dgl.DGLGraph(nx.erdos_renyi_graph(size, 0.3))
+
+
+def random_graph(size):
+    return dgl.from_networkx(nx.erdos_renyi_graph(size, 0.3))
+
+
+def random_bipartite(size_src, size_dst):
+    return dgl.bipartite_from_scipy(
+        ssp.random(size_src, size_dst, 0.1),
+        utype="_U",
+        etype="_E",
+        vtype="V",
+    )
+
+
+def random_block(size):
+    g = dgl.from_networkx(nx.erdos_renyi_graph(size, 0.1))
+    return dgl.to_block(g, np.unique(F.zerocopy_to_numpy(g.edges()[1])))
+
+
+@register_case(["two_hetero_batch"])
+def two_hetero_batch():
+    g1 = dgl.heterograph(
+        {
+            ("user", "follows", "user"): ([0, 1], [1, 2]),
+            ("user", "follows", "developer"): ([0, 1], [1, 2]),
+            ("user", "plays", "game"): ([0, 1, 2, 3], [0, 0, 1, 1]),
+        }
+    )
+    g2 = dgl.heterograph(
+        {
+            ("user", "follows", "user"): ([0, 1], [1, 2]),
+            ("user", "follows", "developer"): ([0, 1], [1, 2]),
+            ("user", "plays", "game"): ([0, 1, 2], [0, 0, 1]),
+        }
+    )
+    return [g1, g2]
+
+
+@register_case(["two_hetero_batch"])
+def two_hetero_batch_with_isolated_ntypes():
+    g1 = dgl.heterograph(
+        {
+            ("user", "follows", "user"): ([0, 1], [1, 2]),
+            ("user", "follows", "developer"): ([0, 1], [1, 2]),
+            ("user", "plays", "game"): ([0, 1, 2, 3], [0, 0, 1, 1]),
+        },
+        num_nodes_dict={"user": 4, "game": 2, "developer": 3, "platform": 2},
+    )
+    g2 = dgl.heterograph(
+        {
+            ("user", "follows", "user"): ([0, 1], [1, 2]),
+            ("user", "follows", "developer"): ([0, 1], [1, 2]),
+            ("user", "plays", "game"): ([0, 1, 2], [0, 0, 1]),
+        },
+        num_nodes_dict={"user": 3, "game": 2, "developer": 3, "platform": 3},
+    )
+    return [g1, g2]
+
+
+@register_case(["batched", "hetero"])
+def batched_heterograph0():
+    g1 = dgl.heterograph(
+        {
+            ("user", "follows", "user"): ([0, 1], [1, 2]),
+            ("user", "follows", "developer"): ([0, 1], [1, 2]),
+            ("user", "plays", "game"): ([0, 1, 2, 3], [0, 0, 1, 1]),
+        }
+    )
+    g2 = dgl.heterograph(
+        {
+            ("user", "follows", "user"): ([0, 1], [1, 2]),
+            ("user", "follows", "developer"): ([0, 1], [1, 2]),
+            ("user", "plays", "game"): ([0, 1, 2], [0, 0, 1]),
+        }
+    )
+    g3 = dgl.heterograph(
+        {
+            ("user", "follows", "user"): ([1], [2]),
+            ("user", "follows", "developer"): ([0, 1, 2], [0, 2, 2]),
+            ("user", "plays", "game"): ([0, 1], [0, 0]),
+        }
+    )
+    return dgl.batch([g1, g2, g3])