()
| 270 | F._default_context_str == "gpu", reason="GPU random walk not implemented" |
| 271 | ) |
| 272 | def test_node2vec(): |
| 273 | g1 = dgl.heterograph({("user", "follow", "user"): ([0, 1, 2], [1, 2, 0])}) |
| 274 | g2 = dgl.heterograph( |
| 275 | {("user", "follow", "user"): ([0, 1, 1, 2, 3], [1, 2, 3, 0, 0])} |
| 276 | ) |
| 277 | g2.edata["p"] = F.tensor([3, 0, 3, 3, 3], dtype=F.float32) |
| 278 | |
| 279 | ntypes = F.zeros((5,), dtype=F.int64) |
| 280 | |
| 281 | traces, eids = dgl.sampling.node2vec_random_walk( |
| 282 | g1, [0, 1, 2, 0, 1, 2], 1, 1, 4, return_eids=True |
| 283 | ) |
| 284 | check_random_walk(g1, ["follow"] * 4, traces, ntypes, trace_eids=eids) |
| 285 | |
| 286 | traces, eids = dgl.sampling.node2vec_random_walk( |
| 287 | g2, [0, 1, 2, 3, 0, 1, 2, 3], 1, 1, 4, prob="p", return_eids=True |
| 288 | ) |
| 289 | check_random_walk(g2, ["follow"] * 4, traces, ntypes, "p", trace_eids=eids) |
| 290 | |
| 291 | |
| 292 | @unittest.skipIf( |
nothing calls this directly
no test coverage detected