(g, bg, ntype, etype, dst_nodes, include_dst_in_src=True)
| 24 | @parametrize_idtype |
| 25 | def test_to_block(idtype): |
| 26 | def check(g, bg, ntype, etype, dst_nodes, include_dst_in_src=True): |
| 27 | if dst_nodes is not None: |
| 28 | assert F.array_equal(bg.dstnodes[ntype].data[dgl.NID], dst_nodes) |
| 29 | n_dst_nodes = bg.num_nodes("DST/" + ntype) |
| 30 | if include_dst_in_src: |
| 31 | assert F.array_equal( |
| 32 | bg.srcnodes[ntype].data[dgl.NID][:n_dst_nodes], |
| 33 | bg.dstnodes[ntype].data[dgl.NID], |
| 34 | ) |
| 35 | |
| 36 | g = g[etype] |
| 37 | bg = bg[etype] |
| 38 | induced_src = bg.srcdata[dgl.NID] |
| 39 | induced_dst = bg.dstdata[dgl.NID] |
| 40 | induced_eid = bg.edata[dgl.EID] |
| 41 | |
| 42 | bg_src, bg_dst = bg.all_edges(order="eid") |
| 43 | src_ans, dst_ans = g.all_edges(order="eid") |
| 44 | |
| 45 | induced_src_bg = F.gather_row(induced_src, bg_src) |
| 46 | induced_dst_bg = F.gather_row(induced_dst, bg_dst) |
| 47 | induced_src_ans = F.gather_row(src_ans, induced_eid) |
| 48 | induced_dst_ans = F.gather_row(dst_ans, induced_eid) |
| 49 | |
| 50 | assert F.array_equal(induced_src_bg, induced_src_ans) |
| 51 | assert F.array_equal(induced_dst_bg, induced_dst_ans) |
| 52 | |
| 53 | def checkall(g, bg, dst_nodes, include_dst_in_src=True): |
| 54 | for etype in g.etypes: |
no test coverage detected