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