Function test_pull_0deg

tests/python/common/function/test_basics.py:434–473

(idtype)

@parametrize_idtype
def test_pull_0deg(idtype):
    g = dgl.graph(([0], [1]), idtype=idtype, device=F.ctx())

    def _message(edges):
        return {"m": edges.src["h"]}

    def _reduce(nodes):
        return {"x": nodes.data["h"] + F.sum(nodes.mailbox["m"], 1)}

    def _apply(nodes):
        return {"x": nodes.data["x"] * 2}

    def _init2(shape, dtype, ctx, ids):
        return 2 + F.zeros(shape, dtype, ctx)

    g.set_n_initializer(_init2, "x")
    # test#1: pull both 0deg and non-0deg nodes
    old = F.randn((2, 5))
    g.ndata["h"] = old
    g.pull([0, 1], _message, _reduce, _apply)
    new = g.ndata["x"]
    # 0deg check: initialized with the func and got applied
    assert F.allclose(new[0], F.full_1d(5, 4, dtype=F.float32))
    # non-0deg check
    assert F.allclose(new[1], F.sum(old, 0) * 2)

    # test#2: pull only 0deg node
    old = F.randn((2, 5))
    g.ndata["h"] = old
    # Intercepting the warning: The input graph for the user-defined edge
    # function does not contain valid edges
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", category=UserWarning)
        g.pull(0, _message, _reduce, lambda nodes: {"h": nodes.data["h"] * 2})

    new = g.ndata["h"]
    # 0deg check: fallback to apply
    assert F.allclose(new[0], 2 * old[0])
    # non-0deg check: not touched
    assert F.allclose(new[1], old[1])


def test_dynamic_addition():