def test_recv_0deg_newfld(): # test recv with 0deg nodes; the reducer also creates a new field g = dgl.graph([(0,1)]) def _message(edges): return {'m' : edges.src['h']} def _reduce(nodes): return {'h1' : nodes.data['h'] + F.sum(nodes.mailbox['m'], 1)} def _apply(nodes): return {'h1' : nodes.data['h1'] * 2} def _init2(shape, dtype, ctx, ids): return 2 + F.zeros(shape, dtype=dtype, ctx=ctx) # test#1: recv both 0deg and non-0deg nodes old = F.randn((2, 5)) g.set_n_initializer(_init2, 'h1') g.ndata['h'] = old g.send((0, 1), _message) g.recv([0, 1], _reduce, _apply) new = g.ndata.pop('h1') # 0deg check: initialized with the func and got applied assert F.allclose(new[0], F.full_1d(5, 4, dtype=F.float32)) # non-0deg check assert F.allclose(new[1], F.sum(old, 0) * 2) # test#2: recv only 0deg node old = F.randn((2, 5)) g.ndata['h'] = old g.ndata['h1'] = F.full((2, 5), -1, F.int64) # this is necessary g.send((0, 1), _message) g.recv(0, _reduce, _apply) new = g.ndata.pop('h1') # 0deg check: fallback to apply assert F.allclose(new[0], F.full_1d(5, -2, F.int64)) # non-0deg check: not changed assert F.allclose(new[1], F.full_1d(5, -1, F.int64))
def test_pull_0deg(index_dtype): g = dgl.graph([(0, 1)], index_dtype=index_dtype) 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, ctx) g.set_n_initializer(_init2, 'h') # test#1: pull both 0deg and non-0deg nodes old = F.randn((2, 5)) g.ndata['h'] = old g.pull([0, 1], _message, _reduce, _apply) new = g.ndata.pop('h') # 0deg check: initialized with the func and got applied assert F.allclose(new[0], F.full_1d(5, 4, dtype=F.float32)) # non-0deg check assert F.allclose(new[1], F.sum(old, 0) * 2) # test#2: pull only 0deg node old = F.randn((2, 5)) g.ndata['h'] = old g.pull(0, _message, _reduce, _apply) new = g.ndata.pop('h') # 0deg check: fallback to apply assert F.allclose(new[0], 2 * old[0]) # non-0deg check: not touched assert F.allclose(new[1], old[1])
def test_recv_0deg(): # test recv with 0deg nodes; g = dgl.graph([(0,1)]) 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, ctx) g.set_n_initializer(_init2, 'h') # test#1: recv both 0deg and non-0deg nodes old = F.randn((2, 5)) g.ndata['h'] = old g.send((0, 1), _message) g.recv([0, 1], _reduce, _apply) new = g.ndata.pop('h') # 0deg check: initialized with the func and got applied assert F.allclose(new[0], F.full_1d(5, 4, F.float32)) # non-0deg check assert F.allclose(new[1], F.sum(old, 0) * 2) # test#2: recv only 0deg node is equal to apply old = F.randn((2, 5)) g.ndata['h'] = old g.send((0, 1), _message) g.recv(0, _reduce, _apply) new = g.ndata.pop('h') # 0deg check: equal to apply_nodes assert F.allclose(new[0], 2 * old[0]) # non-0deg check: untouched assert F.allclose(new[1], old[1])