def test_dynamic_addition():
N = 3
D = 1
g = DGLGraph()
def _init(shape, dtype, ctx, ids):
return F.copy_to(F.astype(F.randn(shape), dtype), ctx)
g.set_n_initializer(_init)
g.set_e_initializer(_init)
def _message(edges):
return {
'm':
edges.src['h1'] + edges.dst['h2'] + edges.data['h1'] +
edges.data['h2']
}
def _reduce(nodes):
return {'h': F.sum(nodes.mailbox['m'], 1)}
def _apply(nodes):
return {'h': nodes.data['h']}
g.register_message_func(_message)
g.register_reduce_func(_reduce)
g.register_apply_node_func(_apply)
g.set_n_initializer(dgl.init.zero_initializer)
g.set_e_initializer(dgl.init.zero_initializer)
# add nodes and edges
g.add_nodes(N)
g.ndata.update({'h1': F.randn((N, D)), 'h2': F.randn((N, D))})
g.add_nodes(3)
g.add_edge(0, 1)
g.add_edge(1, 0)
g.edata.update({'h1': F.randn((2, D)), 'h2': F.randn((2, D))})
g.send()
expected = F.ones((g.number_of_edges(), ), dtype=F.int64)
assert F.array_equal(g._get_msg_index().tousertensor(), expected)
# add more edges
g.add_edges([0, 2], [2, 0], {'h1': F.randn((2, D))})
g.send(([0, 2], [2, 0]))
g.recv(0)
g.add_edge(1, 2)
g.edges[4].data['h1'] = F.randn((1, D))
g.send((1, 2))
g.recv([1, 2])
h = g.ndata.pop('h')
# a complete round of send and recv
g.send()
g.recv()
assert F.allclose(h, g.ndata['h'])
def test_send_recv_after_conversion():
# test send and recv after converting from a graph with edges
g = generate_graph()
# nx graph
nxg = g.to_networkx(node_attrs=['h'])
g1 = DGLGraph()
# some random node and edges
g1.add_nodes(4)
g1.add_edges([1, 2], [2, 3])
g1.set_n_initializer(dgl.init.zero_initializer)
g1.from_networkx(nxg, node_attrs=['h'])
# sparse matrix
row, col = g.all_edges()
data = range(len(row))
n = g.number_of_nodes()
a = sp.coo_matrix(
(data, (F.zerocopy_to_numpy(row), F.zerocopy_to_numpy(col))),
shape=(n, n))
g2 = DGLGraph()
# some random node and edges
g2.add_nodes(5)
g2.add_edges([1, 2, 4], [2, 3, 0])
g2.set_n_initializer(dgl.init.zero_initializer)
g2.from_scipy_sparse_matrix(a)
g2.ndata['h'] = g.ndata['h']
# on dgl graph
g.send(message_func=message_func)
g.recv([0, 1, 3, 5],
reduce_func=reduce_func,
apply_node_func=apply_node_func)
g.recv([0, 2, 4, 8],
reduce_func=reduce_func,
apply_node_func=apply_node_func)
# nx
g1.send(message_func=message_func)
g1.recv([0, 1, 3, 5],
reduce_func=reduce_func,
apply_node_func=apply_node_func)
g1.recv([0, 2, 4, 8],
reduce_func=reduce_func,
apply_node_func=apply_node_func)
# sparse matrix
g2.send(message_func=message_func)
g2.recv([0, 1, 3, 5],
reduce_func=reduce_func,
apply_node_func=apply_node_func)
g2.recv([0, 2, 4, 8],
reduce_func=reduce_func,
apply_node_func=apply_node_func)
assert F.allclose(g.ndata['h'], g1.ndata['h'])
assert F.allclose(g.ndata['h'], g2.ndata['h'])
def check_pull_0deg(readonly):
if readonly:
row_idx = []
col_idx = []
row_idx.append(0)
col_idx.append(1)
ones = np.ones(shape=(len(row_idx)))
csr = spsp.csr_matrix((ones, (row_idx, col_idx)), shape=(2, 2))
g = DGLGraph(csr, readonly=True)
else:
g = DGLGraph()
g.add_nodes(2)
g.add_edge(0, 1)
def _message(edges):
return {'m': edges.src['h']}
def _reduce(nodes):
return {'h': nodes.mailbox['m'].sum(1)}
def _apply(nodes):
return {'h': nodes.data['h'] * 2}
def _init2(shape, dtype, ctx, ids):
return 2 + mx.nd.zeros(shape, dtype=dtype, ctx=ctx)
g.set_n_initializer(_init2, 'h')
old_repr = mx.nd.random.normal(shape=(2, 5))
# test#1: pull only 0-deg node
g.ndata['h'] = old_repr
g.pull(0, _message, _reduce, _apply)
new_repr = g.ndata['h']
# 0deg check: equal to apply_nodes
assert np.allclose(new_repr[0].asnumpy(), old_repr[0].asnumpy() * 2)
# non-0deg check: untouched
assert np.allclose(new_repr[1].asnumpy(), old_repr[1].asnumpy())
# test#2: pull only non-deg node
g.ndata['h'] = old_repr
g.pull(1, _message, _reduce, _apply)
new_repr = g.ndata['h']
# 0deg check: untouched
assert np.allclose(new_repr[0].asnumpy(), old_repr[0].asnumpy())
# non-0deg check: recved node0 and got applied
assert np.allclose(new_repr[1].asnumpy(), old_repr[0].asnumpy() * 2)
# test#3: pull only both nodes
g.ndata['h'] = old_repr
g.pull([0, 1], _message, _reduce, _apply)
new_repr = g.ndata['h']
# 0deg check: init and applied
t = mx.nd.zeros(shape=(2, 5)) + 4
assert np.allclose(new_repr[0].asnumpy(), t.asnumpy())
# non-0deg check: recv node0 and applied
assert np.allclose(new_repr[1].asnumpy(), old_repr[0].asnumpy() * 2)
def generate_graph(grad=False):
g = DGLGraph()
g.add_nodes(10) # 10 nodes.
# create a graph where 0 is the source and 9 is the sink
# 16 edges
for i in range(1, 9):
g.add_edge(0, i)
g.add_edge(i, 9)
ncol = Variable(th.randn(10, D), requires_grad=grad)
ecol = Variable(th.randn(16, D), requires_grad=grad)
g.set_n_initializer(dgl.init.zero_initializer)
g.set_e_initializer(dgl.init.zero_initializer)
g.ndata['h'] = ncol
g.edata['w'] = ecol
return g
def generate_graph(grad=False):
g = DGLGraph()
g.add_nodes(10) # 10 nodes.
# create a graph where 0 is the source and 9 is the sink
# 16 edges
for i in range(1, 9):
g.add_edge(0, i)
g.add_edge(i, 9)
ncol = F.randn((10, D))
ecol = F.randn((16, D))
if grad:
ncol = F.attach_grad(ncol)
ecol = F.attach_grad(ecol)
g.set_n_initializer(dgl.init.zero_initializer)
g.set_e_initializer(dgl.init.zero_initializer)
g.ndata['h'] = ncol
g.edata['w'] = ecol
return g
def test_recv_0deg_newfld():
# test recv with 0deg nodes; the reducer also creates a new field
g = DGLGraph()
g.add_nodes(2)
g.add_edge(0, 1)
def _message(edges):
return {'m': edges.src['h']}
def _reduce(nodes):
return {'h1': nodes.data['h'] + mx.nd.sum(nodes.mailbox['m'], 1)}
def _apply(nodes):
return {'h1': nodes.data['h1'] * 2}
def _init2(shape, dtype, ctx, ids):
return 2 + mx.nd.zeros(shape=shape, dtype=dtype, ctx=ctx)
g.register_message_func(_message)
g.register_reduce_func(_reduce)
g.register_apply_node_func(_apply)
# test#1: recv both 0deg and non-0deg nodes
old = mx.nd.random.normal(shape=(2, 5))
g.set_n_initializer(_init2, 'h1')
g.ndata['h'] = old
g.send((0, 1))
g.recv([0, 1])
new = g.ndata.pop('h1')
# 0deg check: initialized with the func and got applied
assert np.allclose(new[0].asnumpy(), np.full((5, ), 4))
# non-0deg check
assert np.allclose(new[1].asnumpy(), mx.nd.sum(old, 0).asnumpy() * 2)
# test#2: recv only 0deg node
old = mx.nd.random.normal(shape=(2, 5))
g.ndata['h'] = old
g.ndata['h1'] = mx.nd.full((2, 5), -1) # this is necessary
g.send((0, 1))
g.recv(0)
new = g.ndata.pop('h1')
# 0deg check: fallback to apply
assert np.allclose(new[0].asnumpy(), np.full((5, ), -2))
# non-0deg check: not changed
assert np.allclose(new[1].asnumpy(), np.full((5, ), -1))
def generate_graph(grad=False, readonly=False):
if readonly:
row_idx = []
col_idx = []
for i in range(1, 9):
row_idx.append(0)
col_idx.append(i)
row_idx.append(i)
col_idx.append(9)
row_idx.append(9)
col_idx.append(0)
ones = np.ones(shape=(len(row_idx)))
csr = spsp.csr_matrix((ones, (row_idx, col_idx)), shape=(10, 10))
g = DGLGraph(csr, readonly=True)
ncol = mx.nd.random.normal(shape=(10, D))
ecol = mx.nd.random.normal(shape=(17, D))
if grad:
ncol.attach_grad()
ecol.attach_grad()
g.ndata['h'] = ncol
g.edata['w'] = ecol
g.set_n_initializer(dgl.init.zero_initializer)
g.set_e_initializer(dgl.init.zero_initializer)
return g
else:
g = DGLGraph()
g.add_nodes(10) # 10 nodes.
# 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)
ncol = mx.nd.random.normal(shape=(10, D))
ecol = mx.nd.random.normal(shape=(17, D))
if grad:
ncol.attach_grad()
ecol.attach_grad()
g.ndata['h'] = ncol
g.edata['w'] = ecol
g.set_n_initializer(dgl.init.zero_initializer)
g.set_e_initializer(dgl.init.zero_initializer)
return g
def test_recv_0deg():
# test recv with 0deg nodes;
g = DGLGraph()
g.add_nodes(2)
g.add_edge(0, 1)
def _message(edges):
return {'m': edges.src['h']}
def _reduce(nodes):
return {'h': nodes.data['h'] + nodes.mailbox['m'].sum(1)}
def _apply(nodes):
return {'h': nodes.data['h'] * 2}
def _init2(shape, dtype, ctx, ids):
return 2 + th.zeros(shape, dtype=dtype, device=ctx)
g.register_message_func(_message)
g.register_reduce_func(_reduce)
g.register_apply_node_func(_apply)
g.set_n_initializer(_init2, 'h')
# test#1: recv both 0deg and non-0deg nodes
old = th.randn((2, 5))
g.ndata['h'] = old
g.send((0, 1))
g.recv([0, 1])
new = g.ndata.pop('h')
# 0deg check: initialized with the func and got applied
assert U.allclose(new[0], th.full((5, ), 4))
# non-0deg check
assert U.allclose(new[1], th.sum(old, 0) * 2)
# test#2: recv only 0deg node is equal to apply
old = th.randn((2, 5))
g.ndata['h'] = old
g.send((0, 1))
g.recv(0)
new = g.ndata.pop('h')
# 0deg check: equal to apply_nodes
assert U.allclose(new[0], 2 * old[0])
# non-0deg check: untouched
assert U.allclose(new[1], old[1])
def test_update_all_0deg():
# test#1
g = DGLGraph()
g.add_nodes(5)
g.add_edge(1, 0)
g.add_edge(2, 0)
g.add_edge(3, 0)
g.add_edge(4, 0)
def _message(edges):
return {'m': edges.src['h']}
def _reduce(nodes):
return {'h': nodes.data['h'] + mx.nd.sum(nodes.mailbox['m'], 1)}
def _apply(nodes):
return {'h': nodes.data['h'] * 2}
def _init2(shape, dtype, ctx, ids):
return 2 + mx.nd.zeros(shape, dtype=dtype, ctx=ctx)
g.set_n_initializer(_init2, 'h')
old_repr = mx.nd.random.normal(shape=(5, 5))
g.ndata['h'] = old_repr
g.update_all(_message, _reduce, _apply)
new_repr = g.ndata['h']
# the first row of the new_repr should be the sum of all the node
# features; while the 0-deg nodes should be initialized by the
# initializer and applied with UDF.
assert np.allclose(new_repr[1:].asnumpy(), 2 * (2 + np.zeros((4, 5))))
assert np.allclose(new_repr[0].asnumpy(),
2 * mx.nd.sum(old_repr, 0).asnumpy())
# test#2: graph with no edge
g = DGLGraph()
g.add_nodes(5)
g.set_n_initializer(_init2, 'h')
g.ndata['h'] = old_repr
g.update_all(_message, _reduce, _apply)
new_repr = g.ndata['h']
# should fallback to apply
assert np.allclose(new_repr.asnumpy(), 2 * old_repr.asnumpy())
def test_send_twice_different_field():
g = DGLGraph()
g.set_n_initializer(dgl.init.zero_initializer)
g.add_nodes(2)
g.add_edge(0, 1)
def _message_a(edges):
return {'a': edges.src['a']}
def _message_b(edges):
return {'b': edges.src['b']}
def _reduce(nodes):
return {'a': nodes.mailbox['a'].sum(1), 'b': nodes.mailbox['b'].sum(1)}
old_a = th.randn(2, 5)
old_b = th.randn(2, 5)
g.set_n_repr({'a': old_a, 'b': old_b})
g.send((0, 1), _message_a)
g.send((0, 1), _message_b)
g.recv([1], _reduce)
new_repr = g.get_n_repr()
assert th.allclose(new_repr['a'][1], old_a[0])
assert th.allclose(new_repr['b'][1], old_b[0])
def test_multi_recv_0deg():
# test recv with 0deg nodes;
g = DGLGraph()
def _message(edges):
return {'m': edges.src['h']}
def _reduce(nodes):
return {'h': nodes.data['h'] + F.sum(nodes.mailbox['m'], 1)}
def _apply(nodes):
return {'h': nodes.data['h'] * 2}
def _init2(shape, dtype, ctx, ids):
return 2 + F.zeros(shape, dtype=dtype, ctx=ctx)
g.register_message_func(_message)
g.register_reduce_func(_reduce)
g.register_apply_node_func(_apply)
g.set_n_initializer(_init2)
g.add_nodes(2)
g.add_edge(0, 1)
# recv both 0deg and non-0deg nodes
old = F.randn((2, 5))
g.ndata['h'] = old
g.send((0, 1))
g.recv([0, 1])
new = g.ndata['h']
# 0deg check: initialized with the func and got applied
assert F.allclose(new[0], F.full((5, ), 4, F.float32))
# non-0deg check
assert F.allclose(new[1], F.sum(old, 0) * 2)
# recv again on zero degree node
g.recv([0])
assert F.allclose(g.nodes[0].data['h'], F.full((5, ), 8, F.float32))
# recv again on node with no incoming message
g.recv([1])
assert F.allclose(g.nodes[1].data['h'], F.sum(old, 0) * 4)