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.copy_to(F.ones((g.number_of_edges(), ), dtype=F.int64),
F.cpu())
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 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 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