def check_reduce_0deg(readonly):
if readonly:
row_idx = []
col_idx = []
for i in range(1, 5):
row_idx.append(i)
col_idx.append(0)
ones = np.ones(shape=(len(row_idx)))
csr = spsp.csr_matrix((ones, (row_idx, col_idx)), shape=(5, 5))
g = DGLGraph(csr, readonly=True)
else:
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'] + nodes.mailbox['m'].sum(1)}
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.set_n_repr({'h': old_repr})
g.update_all(_message, _reduce)
new_repr = g.ndata['h']
assert np.allclose(new_repr[1:].asnumpy(), 2 + np.zeros((4, 5)))
assert np.allclose(new_repr[0].asnumpy(), old_repr.sum(0).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': F.sum(nodes.mailbox['a'], 1),
'b': F.sum(nodes.mailbox['b'], 1)
}
old_a = F.randn((2, 5))
old_b = F.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 F.allclose(new_repr['a'][1], old_a[0])
assert F.allclose(new_repr['b'][1], old_b[0])
