Function tree2

tests/python/common/test_batch-graph.py:30–47

Generate a tree 1 / \ 4 3 / \ 2 0 Edges are from leaves to root.

(idtype)

def tree2(idtype):
    """Generate a tree
         1
        / \
       4   3
      / \
     2   0
    Edges are from leaves to root.
    """
    g = dgl.graph(([], [])).astype(idtype).to(F.ctx())
    g.add_nodes(5)
    g.add_edges(2, 4)
    g.add_edges(0, 4)
    g.add_edges(4, 1)
    g.add_edges(3, 1)
    g.ndata["h"] = F.tensor([0, 1, 2, 3, 4])
    g.edata["h"] = F.randn((4, 10))
    return g


@parametrize_idtype