(max_degree, embedding_dim, direction)
| 2452 | @pytest.mark.parametrize("embedding_dim", [8, 16]) |
| 2453 | @pytest.mark.parametrize("direction", ["in", "out", "both"]) |
| 2454 | def test_degree_encoder(max_degree, embedding_dim, direction): |
| 2455 | g1 = dgl.graph( |
| 2456 | ( |
| 2457 | th.tensor([0, 0, 0, 1, 1, 2, 3, 3]), |
| 2458 | th.tensor([1, 2, 3, 0, 3, 0, 0, 1]), |
| 2459 | ) |
| 2460 | ) |
| 2461 | g2 = dgl.graph( |
| 2462 | ( |
| 2463 | th.tensor([0, 1]), |
| 2464 | th.tensor([1, 0]), |
| 2465 | ) |
| 2466 | ) |
| 2467 | in_degree = pad_sequence( |
| 2468 | [g1.in_degrees(), g2.in_degrees()], batch_first=True |
| 2469 | ) |
| 2470 | out_degree = pad_sequence( |
| 2471 | [g1.out_degrees(), g2.out_degrees()], batch_first=True |
| 2472 | ) |
| 2473 | model = nn.DegreeEncoder(max_degree, embedding_dim, direction=direction) |
| 2474 | if direction == "in": |
| 2475 | de_g = model(in_degree) |
| 2476 | elif direction == "out": |
| 2477 | de_g = model(out_degree) |
| 2478 | elif direction == "both": |
| 2479 | de_g = model(th.stack((in_degree, out_degree))) |
| 2480 | assert de_g.shape == (2, 4, embedding_dim) |
| 2481 | |
| 2482 | |
| 2483 | @parametrize_idtype |
nothing calls this directly
no test coverage detected