| 2722 | |
| 2723 | @parametrize_idtype |
| 2724 | def test_module_khop_graph(idtype): |
| 2725 | transform = dgl.KHopGraph(2) |
| 2726 | g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx()) |
| 2727 | g.ndata["h"] = F.randn((g.num_nodes(), 2)) |
| 2728 | new_g = transform(g) |
| 2729 | assert new_g.device == g.device |
| 2730 | assert new_g.idtype == g.idtype |
| 2731 | assert new_g.num_nodes() == g.num_nodes() |
| 2732 | assert F.allclose(g.ndata["h"], new_g.ndata["h"]) |
| 2733 | src, dst = new_g.edges() |
| 2734 | eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst)))) |
| 2735 | assert eset == {(0, 2)} |
| 2736 | |
| 2737 | |
| 2738 | @parametrize_idtype |