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])