"vscode:/vscode.git/clone" did not exist on "e6d46123091afd58281dc7487c0f6b67055683b9"
test_basics.py 1.88 KB
Newer Older
1
2
# NOTE(vibwu): Currently cugraph must be imported before torch to avoid a resource cleanup issue.
#    See https://github.com/rapidsai/cugraph/issues/2718
3
import cugraph  # usort: skip
4
import backend as F
5

6
7
import dgl

8
9
10
11

def test_dummy():
    cg = cugraph.Graph()
    assert cg is not None
12

13

14
def test_to_cugraph_conversion():
15
    g = dgl.graph((F.tensor([0, 1, 2, 3]), F.tensor([1, 0, 3, 2]))).to("cuda")
16
17
    cugraph_g = g.to_cugraph()

Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
18
19
    assert cugraph_g.number_of_nodes() == g.num_nodes()
    assert cugraph_g.number_of_edges() == g.num_edges()
20
21
22
23
24

    assert cugraph_g.has_edge(0, 1)
    assert cugraph_g.has_edge(1, 0)
    assert cugraph_g.has_edge(3, 2)

25

26
27
28
29
30
31
def test_from_cugraph_conversion():
    # cudf is a dependency of cugraph
    import cudf

    # directed graph conversion test
    cugraph_g = cugraph.Graph(directed=True)
32
    df = cudf.DataFrame({"source": [0, 1, 2, 3], "destination": [1, 2, 3, 2]})
33
34
35
36
37

    cugraph_g.from_cudf_edgelist(df)

    g = dgl.from_cugraph(cugraph_g)

38
    assert g.device.type == "cuda"
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
39
40
    assert g.num_nodes() == cugraph_g.number_of_nodes()
    assert g.num_edges() == cugraph_g.number_of_edges()
41
42
43
44
45
46
47
48
49
50

    # assert reverse edges are not present
    assert g.has_edges_between(0, 1)
    assert not g.has_edges_between(1, 0)
    assert g.has_edges_between(1, 2)
    assert not g.has_edges_between(2, 1)
    assert g.has_edges_between(2, 3)

    # undirected graph conversion test
    cugraph_g = cugraph.Graph(directed=False)
51
    df = cudf.DataFrame({"source": [0, 1, 2, 3], "destination": [1, 2, 3, 2]})
52
53
54
55
56

    cugraph_g.from_cudf_edgelist(df)

    g = dgl.from_cugraph(cugraph_g)

57
    assert g.device.type == "cuda"
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
58
    assert g.num_nodes() == cugraph_g.number_of_nodes()
59
60
61
62
63
64
    # assert reverse edges are present
    assert g.has_edges_between(0, 1)
    assert g.has_edges_between(1, 0)
    assert g.has_edges_between(1, 2)
    assert g.has_edges_between(2, 1)
    assert g.has_edges_between(2, 3)