Function test_batch_propagate

tests/python/common/test_batch-graph.py:167–192

(idtype)

@parametrize_idtype
def test_batch_propagate(idtype):
    t1 = tree1(idtype)
    t2 = tree2(idtype)

    bg = dgl.batch([t1, t2])
    _mfunc = lambda edges: {"m": edges.src["h"]}
    _rfunc = lambda nodes: {"h": F.sum(nodes.mailbox["m"], 1)}
    # get leaves.

    order = []

    # step 1
    u = [3, 4, 2 + 5, 0 + 5]
    v = [1, 1, 4 + 5, 4 + 5]
    order.append((u, v))

    # step 2
    u = [1, 2, 4 + 5, 3 + 5]
    v = [0, 0, 1 + 5, 1 + 5]
    order.append((u, v))

    bg.prop_edges(order, _mfunc, _rfunc)
    t1, t2 = dgl.unbatch(bg)

    assert F.asnumpy(t1.ndata["h"][0]) == 9
    assert F.asnumpy(t2.ndata["h"][1]) == 5


@parametrize_idtype