Function test_level2

tests/python/common/test_heterograph.py:1845–1978

(idtype)

@parametrize_idtype
def test_level2(idtype):
    # edges = {
    #    'follows': ([0, 1], [1, 2]),
    #    'plays': ([0, 1, 2, 1], [0, 0, 1, 1]),
    #    'wishes': ([0, 2], [1, 0]),
    #    'develops': ([0, 1], [0, 1]),
    # }
    g = create_test_heterograph(idtype)

    def rfunc(nodes):
        return {"y": F.sum(nodes.mailbox["m"], 1)}

    def rfunc2(nodes):
        return {"y": F.max(nodes.mailbox["m"], 1)}

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

    def afunc(nodes):
        return {"y": nodes.data["y"] + 1}

    #############################################################
    #  send_and_recv
    #############################################################

    g.nodes["user"].data["h"] = F.ones((3, 2))
    g.send_and_recv([2, 3], mfunc, rfunc, etype="plays")
    y = g.nodes["game"].data["y"]
    assert F.array_equal(y, F.tensor([[0.0, 0.0], [2.0, 2.0]]))

    # only one type
    g["plays"].send_and_recv([2, 3], mfunc, rfunc)
    y = g.nodes["game"].data["y"]
    assert F.array_equal(y, F.tensor([[0.0, 0.0], [2.0, 2.0]]))

    # test fail case
    # fail due to multiple types
    with pytest.raises(DGLError):
        g.send_and_recv([2, 3], mfunc, rfunc)

    g.nodes["game"].data.clear()

    #############################################################
    #  pull
    #############################################################

    g.nodes["user"].data["h"] = F.ones((3, 2))
    g.pull(1, mfunc, rfunc, etype="plays")
    y = g.nodes["game"].data["y"]
    assert F.array_equal(y, F.tensor([[0.0, 0.0], [2.0, 2.0]]))

    # only one type
    g["plays"].pull(1, mfunc, rfunc)
    y = g.nodes["game"].data["y"]
    assert F.array_equal(y, F.tensor([[0.0, 0.0], [2.0, 2.0]]))

    # test fail case
    with pytest.raises(DGLError):