Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
test_function.py 5.49 KB
Newer Older
Minjie Wang's avatar
Builtin function and API changes (#53)  
Minjie Wang committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159 
import torch as th
import dgl
import dgl.function as fn
from dgl.graph import __REPR__

def generate_graph():
    g = dgl.DGLGraph()
    for i in range(10):
        g.add_node(i) # 10 nodes.
    h = th.arange(1, 11)
    g.set_n_repr({'h': h})
    # create a graph where 0 is the source and 9 is the sink
    for i in range(1, 9):
        g.add_edge(0, i)
        g.add_edge(i, 9)
    # add a back flow from 9 to 0
    g.add_edge(9, 0)
    h = th.tensor([1., 2., 1., 3., 1., 4., 1., 5., 1., 6.,\
            1., 7., 1., 8., 1., 9., 10.])
    g.set_e_repr({'h' : h})
    return g

def generate_graph1():
    """graph with anonymous repr"""
    g = dgl.DGLGraph()
    for i in range(10):
        g.add_node(i) # 10 nodes.
    h = th.arange(1, 11)
    g.set_n_repr(h)
    # create a graph where 0 is the source and 9 is the sink
    for i in range(1, 9):
        g.add_edge(0, i)
        g.add_edge(i, 9)
    # add a back flow from 9 to 0
    g.add_edge(9, 0)
    h = th.tensor([1., 2., 1., 3., 1., 4., 1., 5., 1., 6.,\
            1., 7., 1., 8., 1., 9., 10.])
    g.set_e_repr(h)
    return g

def reducer_msg(node, msgs):
    return th.sum(msgs['m'], 1)

def reducer_out(node, msgs):
    return {'h' : th.sum(msgs, 1)}

def reducer_both(node, msgs):
    return {'h' : th.sum(msgs['m'], 1)}

def reducer_none(node, msgs):
    return th.sum(msgs, 1)

def test_copy_src():
    # copy_src with both fields
    g = generate_graph()
    g.register_message_func(fn.copy_src(src='h', out='m'), batchable=True)
    g.register_reduce_func(reducer_both, batchable=True)
    g.update_all()
    assert th.allclose(g.get_n_repr()['h'],
            th.tensor([10., 1., 1., 1., 1., 1., 1., 1., 1., 44.]))

    # copy_src with only src field; the out field should use anonymous repr
    g = generate_graph()
    g.register_message_func(fn.copy_src(src='h'), batchable=True)
    g.register_reduce_func(reducer_out, batchable=True)
    g.update_all()
    assert th.allclose(g.get_n_repr()['h'],
            th.tensor([10., 1., 1., 1., 1., 1., 1., 1., 1., 44.]))

    # copy_src with no src field; should use anonymous repr
    g = generate_graph1()
    g.register_message_func(fn.copy_src(out='m'), batchable=True)
    g.register_reduce_func(reducer_both, batchable=True)
    g.update_all()
    assert th.allclose(g.get_n_repr()['h'],
            th.tensor([10., 1., 1., 1., 1., 1., 1., 1., 1., 44.]))

    # copy src with no fields;
    g = generate_graph1()
    g.register_message_func(fn.copy_src(), batchable=True)
    g.register_reduce_func(reducer_out, batchable=True)
    g.update_all()
    assert th.allclose(g.get_n_repr()['h'],
            th.tensor([10., 1., 1., 1., 1., 1., 1., 1., 1., 44.]))

def test_copy_edge():
    # copy_edge with both fields
    g = generate_graph()
    g.register_message_func(fn.copy_edge(edge='h', out='m'), batchable=True)
    g.register_reduce_func(reducer_both, batchable=True)
    g.update_all()
    assert th.allclose(g.get_n_repr()['h'],
            th.tensor([10., 1., 1., 1., 1., 1., 1., 1., 1., 44.]))

    # copy_edge with only edge field; the out field should use anonymous repr
    g = generate_graph()
    g.register_message_func(fn.copy_edge(edge='h'), batchable=True)
    g.register_reduce_func(reducer_out, batchable=True)
    g.update_all()
    assert th.allclose(g.get_n_repr()['h'],
            th.tensor([10., 1., 1., 1., 1., 1., 1., 1., 1., 44.]))

    # copy_edge with no edge field; should use anonymous repr
    g = generate_graph1()
    g.register_message_func(fn.copy_edge(out='m'), batchable=True)
    g.register_reduce_func(reducer_both, batchable=True)
    g.update_all()
    assert th.allclose(g.get_n_repr()['h'],
            th.tensor([10., 1., 1., 1., 1., 1., 1., 1., 1., 44.]))

    # copy edge with no fields;
    g = generate_graph1()
    g.register_message_func(fn.copy_edge(), batchable=True)
    g.register_reduce_func(reducer_out, batchable=True)
    g.update_all()
    assert th.allclose(g.get_n_repr()['h'],
            th.tensor([10., 1., 1., 1., 1., 1., 1., 1., 1., 44.]))

def test_src_mul_edge():
    # src_mul_edge with all fields
    g = generate_graph()
    g.register_message_func(fn.src_mul_edge(src='h', edge='h', out='m'), batchable=True)
    g.register_reduce_func(reducer_both, batchable=True)
    g.update_all()
    assert th.allclose(g.get_n_repr()['h'],
            th.tensor([100., 1., 1., 1., 1., 1., 1., 1., 1., 284.]))

    g = generate_graph()
    g.register_message_func(fn.src_mul_edge(src='h', edge='h'), batchable=True)
    g.register_reduce_func(reducer_out, batchable=True)
    g.update_all()
    assert th.allclose(g.get_n_repr()['h'],
            th.tensor([100., 1., 1., 1., 1., 1., 1., 1., 1., 284.]))

    g = generate_graph1()
    g.register_message_func(fn.src_mul_edge(out='m'), batchable=True)
    g.register_reduce_func(reducer_both, batchable=True)
    g.update_all()
    assert th.allclose(g.get_n_repr()['h'],
            th.tensor([100., 1., 1., 1., 1., 1., 1., 1., 1., 284.]))

    g = generate_graph1()
    g.register_message_func(fn.src_mul_edge(), batchable=True)
    g.register_reduce_func(reducer_out, batchable=True)
    g.update_all()
    assert th.allclose(g.get_n_repr()['h'],
            th.tensor([100., 1., 1., 1., 1., 1., 1., 1., 1., 284.]))

    g = generate_graph1()
    g.register_message_func(fn.src_mul_edge(), batchable=True)
    g.register_reduce_func(reducer_none, batchable=True)
    g.update_all()
    assert th.allclose(g.get_n_repr(),
            th.tensor([100., 1., 1., 1., 1., 1., 1., 1., 1., 284.]))

if __name__ == '__main__':
    test_copy_src()
    test_copy_edge()
    test_src_mul_edge()