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.ones((g.number_of_edges(), ), dtype=F.int64) 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 test_send_recv_after_conversion(): # test send and recv after converting from a graph with edges g = generate_graph() # nx graph nxg = g.to_networkx(node_attrs=['h']) g1 = DGLGraph() # some random node and edges g1.add_nodes(4) g1.add_edges([1, 2], [2, 3]) g1.set_n_initializer(dgl.init.zero_initializer) g1.from_networkx(nxg, node_attrs=['h']) # sparse matrix row, col = g.all_edges() data = range(len(row)) n = g.number_of_nodes() a = sp.coo_matrix( (data, (F.zerocopy_to_numpy(row), F.zerocopy_to_numpy(col))), shape=(n, n)) g2 = DGLGraph() # some random node and edges g2.add_nodes(5) g2.add_edges([1, 2, 4], [2, 3, 0]) g2.set_n_initializer(dgl.init.zero_initializer) g2.from_scipy_sparse_matrix(a) g2.ndata['h'] = g.ndata['h'] # on dgl graph g.send(message_func=message_func) g.recv([0, 1, 3, 5], reduce_func=reduce_func, apply_node_func=apply_node_func) g.recv([0, 2, 4, 8], reduce_func=reduce_func, apply_node_func=apply_node_func) # nx g1.send(message_func=message_func) g1.recv([0, 1, 3, 5], reduce_func=reduce_func, apply_node_func=apply_node_func) g1.recv([0, 2, 4, 8], reduce_func=reduce_func, apply_node_func=apply_node_func) # sparse matrix g2.send(message_func=message_func) g2.recv([0, 1, 3, 5], reduce_func=reduce_func, apply_node_func=apply_node_func) g2.recv([0, 2, 4, 8], reduce_func=reduce_func, apply_node_func=apply_node_func) assert F.allclose(g.ndata['h'], g1.ndata['h']) assert F.allclose(g.ndata['h'], g2.ndata['h'])
def check_pull_0deg(readonly): if readonly: row_idx = [] col_idx = [] row_idx.append(0) col_idx.append(1) ones = np.ones(shape=(len(row_idx))) csr = spsp.csr_matrix((ones, (row_idx, col_idx)), shape=(2, 2)) g = DGLGraph(csr, readonly=True) else: g = DGLGraph() g.add_nodes(2) g.add_edge(0, 1) def _message(edges): return {'m': edges.src['h']} def _reduce(nodes): return {'h': nodes.mailbox['m'].sum(1)} def _apply(nodes): return {'h': nodes.data['h'] * 2} 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=(2, 5)) # test#1: pull only 0-deg node g.ndata['h'] = old_repr g.pull(0, _message, _reduce, _apply) new_repr = g.ndata['h'] # 0deg check: equal to apply_nodes assert np.allclose(new_repr[0].asnumpy(), old_repr[0].asnumpy() * 2) # non-0deg check: untouched assert np.allclose(new_repr[1].asnumpy(), old_repr[1].asnumpy()) # test#2: pull only non-deg node g.ndata['h'] = old_repr g.pull(1, _message, _reduce, _apply) new_repr = g.ndata['h'] # 0deg check: untouched assert np.allclose(new_repr[0].asnumpy(), old_repr[0].asnumpy()) # non-0deg check: recved node0 and got applied assert np.allclose(new_repr[1].asnumpy(), old_repr[0].asnumpy() * 2) # test#3: pull only both nodes g.ndata['h'] = old_repr g.pull([0, 1], _message, _reduce, _apply) new_repr = g.ndata['h'] # 0deg check: init and applied t = mx.nd.zeros(shape=(2, 5)) + 4 assert np.allclose(new_repr[0].asnumpy(), t.asnumpy()) # non-0deg check: recv node0 and applied assert np.allclose(new_repr[1].asnumpy(), old_repr[0].asnumpy() * 2)
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 test_recv_0deg_newfld(): # test recv with 0deg nodes; the reducer also creates a new field g = DGLGraph() g.add_nodes(2) g.add_edge(0, 1) def _message(edges): return {'m': edges.src['h']} def _reduce(nodes): return {'h1': nodes.data['h'] + mx.nd.sum(nodes.mailbox['m'], 1)} def _apply(nodes): return {'h1': nodes.data['h1'] * 2} def _init2(shape, dtype, ctx, ids): return 2 + mx.nd.zeros(shape=shape, dtype=dtype, ctx=ctx) g.register_message_func(_message) g.register_reduce_func(_reduce) g.register_apply_node_func(_apply) # test#1: recv both 0deg and non-0deg nodes old = mx.nd.random.normal(shape=(2, 5)) g.set_n_initializer(_init2, 'h1') g.ndata['h'] = old g.send((0, 1)) g.recv([0, 1]) new = g.ndata.pop('h1') # 0deg check: initialized with the func and got applied assert np.allclose(new[0].asnumpy(), np.full((5, ), 4)) # non-0deg check assert np.allclose(new[1].asnumpy(), mx.nd.sum(old, 0).asnumpy() * 2) # test#2: recv only 0deg node old = mx.nd.random.normal(shape=(2, 5)) g.ndata['h'] = old g.ndata['h1'] = mx.nd.full((2, 5), -1) # this is necessary g.send((0, 1)) g.recv(0) new = g.ndata.pop('h1') # 0deg check: fallback to apply assert np.allclose(new[0].asnumpy(), np.full((5, ), -2)) # non-0deg check: not changed assert np.allclose(new[1].asnumpy(), np.full((5, ), -1))
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
def test_recv_0deg(): # test recv with 0deg nodes; g = DGLGraph() g.add_nodes(2) g.add_edge(0, 1) def _message(edges): return {'m': edges.src['h']} def _reduce(nodes): return {'h': nodes.data['h'] + nodes.mailbox['m'].sum(1)} def _apply(nodes): return {'h': nodes.data['h'] * 2} def _init2(shape, dtype, ctx, ids): return 2 + th.zeros(shape, dtype=dtype, device=ctx) g.register_message_func(_message) g.register_reduce_func(_reduce) g.register_apply_node_func(_apply) g.set_n_initializer(_init2, 'h') # test#1: recv both 0deg and non-0deg nodes old = th.randn((2, 5)) g.ndata['h'] = old g.send((0, 1)) g.recv([0, 1]) new = g.ndata.pop('h') # 0deg check: initialized with the func and got applied assert U.allclose(new[0], th.full((5, ), 4)) # non-0deg check assert U.allclose(new[1], th.sum(old, 0) * 2) # test#2: recv only 0deg node is equal to apply old = th.randn((2, 5)) g.ndata['h'] = old g.send((0, 1)) g.recv(0) new = g.ndata.pop('h') # 0deg check: equal to apply_nodes assert U.allclose(new[0], 2 * old[0]) # non-0deg check: untouched assert U.allclose(new[1], old[1])
def test_update_all_0deg(): # test#1 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'] + mx.nd.sum(nodes.mailbox['m'], 1)} def _apply(nodes): return {'h': nodes.data['h'] * 2} 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.ndata['h'] = old_repr g.update_all(_message, _reduce, _apply) new_repr = g.ndata['h'] # the first row of the new_repr should be the sum of all the node # features; while the 0-deg nodes should be initialized by the # initializer and applied with UDF. assert np.allclose(new_repr[1:].asnumpy(), 2 * (2 + np.zeros((4, 5)))) assert np.allclose(new_repr[0].asnumpy(), 2 * mx.nd.sum(old_repr, 0).asnumpy()) # test#2: graph with no edge g = DGLGraph() g.add_nodes(5) g.set_n_initializer(_init2, 'h') g.ndata['h'] = old_repr g.update_all(_message, _reduce, _apply) new_repr = g.ndata['h'] # should fallback to apply assert np.allclose(new_repr.asnumpy(), 2 * old_repr.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': nodes.mailbox['a'].sum(1), 'b': nodes.mailbox['b'].sum(1)} old_a = th.randn(2, 5) old_b = th.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 th.allclose(new_repr['a'][1], old_a[0]) assert th.allclose(new_repr['b'][1], old_b[0])
def test_multi_recv_0deg(): # test recv with 0deg nodes; g = DGLGraph() 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=dtype, ctx=ctx) g.register_message_func(_message) g.register_reduce_func(_reduce) g.register_apply_node_func(_apply) g.set_n_initializer(_init2) g.add_nodes(2) g.add_edge(0, 1) # recv both 0deg and non-0deg nodes old = F.randn((2, 5)) g.ndata['h'] = old g.send((0, 1)) g.recv([0, 1]) new = g.ndata['h'] # 0deg check: initialized with the func and got applied assert F.allclose(new[0], F.full((5, ), 4, F.float32)) # non-0deg check assert F.allclose(new[1], F.sum(old, 0) * 2) # recv again on zero degree node g.recv([0]) assert F.allclose(g.nodes[0].data['h'], F.full((5, ), 8, F.float32)) # recv again on node with no incoming message g.recv([1]) assert F.allclose(g.nodes[1].data['h'], F.sum(old, 0) * 4)