| 2815 | |
| 2816 | @parametrize_idtype |
| 2817 | def test_module_compose(idtype): |
| 2818 | g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx()) |
| 2819 | transform = dgl.Compose([dgl.AddReverse(), dgl.AddSelfLoop()]) |
| 2820 | new_g = transform(g) |
| 2821 | assert new_g.device == g.device |
| 2822 | assert new_g.idtype == g.idtype |
| 2823 | assert new_g.num_edges() == 7 |
| 2824 | |
| 2825 | src, dst = new_g.edges() |
| 2826 | eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst)))) |
| 2827 | assert eset == {(0, 1), (1, 2), (1, 0), (2, 1), (0, 0), (1, 1), (2, 2)} |
| 2828 | |
| 2829 | |
| 2830 | @parametrize_idtype |