def test_pickling_batched_heterograph(): # copied from test_heterograph.create_test_heterograph() plays_spmat = ssp.coo_matrix(([1, 1, 1, 1], ([0, 1, 2, 1], [0, 0, 1, 1]))) wishes_nx = nx.DiGraph() wishes_nx.add_nodes_from(['u0', 'u1', 'u2'], bipartite=0) wishes_nx.add_nodes_from(['g0', 'g1'], bipartite=1) wishes_nx.add_edge('u0', 'g1', id=0) wishes_nx.add_edge('u2', 'g0', id=1) follows_g = dgl.graph([(0, 1), (1, 2)], 'user', 'follows') plays_g = dgl.bipartite(plays_spmat, 'user', 'plays', 'game') wishes_g = dgl.bipartite(wishes_nx, 'user', 'wishes', 'game') develops_g = dgl.bipartite([(0, 0), (1, 1)], 'developer', 'develops', 'game') g = dgl.hetero_from_relations([follows_g, plays_g, wishes_g, develops_g]) g2 = dgl.hetero_from_relations([follows_g, plays_g, wishes_g, develops_g]) g.nodes['user'].data['u_h'] = F.randn((3, 4)) g.nodes['game'].data['g_h'] = F.randn((2, 5)) g.edges['plays'].data['p_h'] = F.randn((4, 6)) g2.nodes['user'].data['u_h'] = F.randn((3, 4)) g2.nodes['game'].data['g_h'] = F.randn((2, 5)) g2.edges['plays'].data['p_h'] = F.randn((4, 6)) bg = dgl.batch_hetero([g, g2]) new_bg = _reconstruct_pickle(bg) test_utils.check_graph_equal(bg, new_bg)
def test_pickling_batched_heterograph(): # copied from test_heterograph.create_test_heterograph() g = dgl.heterograph({ ('user', 'follows', 'user'): ([0, 1], [1, 2]), ('user', 'plays', 'game'): ([0, 1, 2, 1], [0, 0, 1, 1]), ('user', 'wishes', 'game'): ([0, 2], [1, 0]), ('developer', 'develops', 'game'): ([0, 1], [0, 1]) }) g2 = dgl.heterograph({ ('user', 'follows', 'user'): ([0, 1], [1, 2]), ('user', 'plays', 'game'): ([0, 1, 2, 1], [0, 0, 1, 1]), ('user', 'wishes', 'game'): ([0, 2], [1, 0]), ('developer', 'develops', 'game'): ([0, 1], [0, 1]) }) g.nodes['user'].data['u_h'] = F.randn((3, 4)) g.nodes['game'].data['g_h'] = F.randn((2, 5)) g.edges['plays'].data['p_h'] = F.randn((4, 6)) g2.nodes['user'].data['u_h'] = F.randn((3, 4)) g2.nodes['game'].data['g_h'] = F.randn((2, 5)) g2.edges['plays'].data['p_h'] = F.randn((4, 6)) bg = dgl.batch_hetero([g, g2]) new_bg = _reconstruct_pickle(bg) test_utils.check_graph_equal(bg, new_bg)
def test_unbatch2(idtype): # batch 3 graphs but unbatch to 2 g1 = dgl.graph(([0, 1, 2], [1, 2, 3]), idtype=idtype, device=F.ctx()) g2 = dgl.graph(([0, 1, 2], [1, 2, 3]), idtype=idtype, device=F.ctx()) g3 = dgl.graph(([0, 1, 2], [1, 2, 3]), idtype=idtype, device=F.ctx()) bg = dgl.batch([g1, g2, g3]) bnn = F.tensor([8, 4]) bne = F.tensor([6, 3]) f1, f2 = dgl.unbatch(bg, node_split=bnn, edge_split=bne) u, v = f1.edges(order='eid') assert F.allclose(u, F.tensor([0, 1, 2, 4, 5, 6])) assert F.allclose(v, F.tensor([1, 2, 3, 5, 6, 7])) u, v = f2.edges(order='eid') assert F.allclose(u, F.tensor([0, 1, 2])) assert F.allclose(v, F.tensor([1, 2, 3])) # batch 2 but unbatch to 3 bg = dgl.batch([f1, f2]) gg1, gg2, gg3 = dgl.unbatch(bg, F.tensor([4, 4, 4]), F.tensor([3, 3, 3])) check_graph_equal(g1, gg1) check_graph_equal(g2, gg2) check_graph_equal(g3, gg3)
def test_pickling_graph(g, idtype): g = g.astype(idtype) new_g = _reconstruct_pickle(g) test_utils.check_graph_equal(g, new_g, check_feature=True)