(grad=False, add_data=True)
| 13 | |
| 14 | |
| 15 | def generate_graph(grad=False, add_data=True): |
| 16 | g = dgl.graph([]).to(F.ctx()) |
| 17 | g.add_nodes(10) |
| 18 | # create a graph where 0 is the source and 9 is the sink |
| 19 | for i in range(1, 9): |
| 20 | g.add_edges(0, i) |
| 21 | g.add_edges(i, 9) |
| 22 | # add a back flow from 9 to 0 |
| 23 | g.add_edges(9, 0) |
| 24 | if add_data: |
| 25 | ncol = F.randn((10, D)) |
| 26 | ecol = F.randn((17, D)) |
| 27 | if grad: |
| 28 | ncol = F.attach_grad(ncol) |
| 29 | ecol = F.attach_grad(ecol) |
| 30 | g.ndata["h"] = ncol |
| 31 | g.edata["l"] = ecol |
| 32 | return g |
| 33 | |
| 34 | |
| 35 | def test_edge_subgraph(): |