| 25 | |
| 26 | |
| 27 | def generate_graph(idtype, grad=False, add_data=True): |
| 28 | g = dgl.graph([]).to(F.ctx(), dtype=idtype) |
| 29 | g.add_nodes(10) |
| 30 | u, v = [], [] |
| 31 | # create a graph where 0 is the source and 9 is the sink |
| 32 | for i in range(1, 9): |
| 33 | u.append(0) |
| 34 | v.append(i) |
| 35 | u.append(i) |
| 36 | v.append(9) |
| 37 | # add a back flow from 9 to 0 |
| 38 | u.append(9) |
| 39 | v.append(0) |
| 40 | g.add_edges(u, v) |
| 41 | if add_data: |
| 42 | ncol = F.randn((10, D)) |
| 43 | ecol = F.randn((17, D)) |
| 44 | if grad: |
| 45 | ncol = F.attach_grad(ncol) |
| 46 | ecol = F.attach_grad(ecol) |
| 47 | g.ndata["h"] = ncol |
| 48 | g.edata["l"] = ecol |
| 49 | return g |
| 50 | |
| 51 | |
| 52 | @unittest.skipIf(not F.gpu_ctx(), reason="only necessary with GPU") |