Function test_update_routines

tests/python/common/function/test_basics.py:352–386

(idtype)

@parametrize_idtype
def test_update_routines(idtype):
    g = generate_graph(idtype)

    # send_and_recv
    reduce_msg_shapes.clear()
    u = [0, 0, 0, 4, 5, 6]
    v = [1, 2, 3, 9, 9, 9]
    g.send_and_recv((u, v), message_func, reduce_func, apply_node_func)
    assert reduce_msg_shapes == {(1, 3, D), (3, 1, D)}
    reduce_msg_shapes.clear()
    try:
        g.send_and_recv([u, v])
        assert False
    except:
        pass

    # pull
    v = F.tensor([1, 2, 3, 9], g.idtype)
    reduce_msg_shapes.clear()
    g.pull(v, message_func, reduce_func, apply_node_func)
    assert reduce_msg_shapes == {(1, 8, D), (3, 1, D)}
    reduce_msg_shapes.clear()

    # push
    v = F.tensor([0, 1, 2, 3], g.idtype)
    reduce_msg_shapes.clear()
    g.push(v, message_func, reduce_func, apply_node_func)
    assert reduce_msg_shapes == {(1, 3, D), (8, 1, D)}
    reduce_msg_shapes.clear()

    # update_all
    reduce_msg_shapes.clear()
    g.update_all(message_func, reduce_func, apply_node_func)
    assert reduce_msg_shapes == {(1, 8, D), (9, 1, D)}
    reduce_msg_shapes.clear()


@parametrize_idtype