Function test_batch_setter_getter

tests/python/common/function/test_basics.py:100–179

(idtype)

@parametrize_idtype
def test_batch_setter_getter(idtype):
    def _pfc(x):
        return list(F.zerocopy_to_numpy(x)[:, 0])

    g = generate_graph(idtype)
    # set all nodes
    g.ndata["h"] = F.zeros((10, D))
    assert F.allclose(g.ndata["h"], F.zeros((10, D)))
    # pop nodes
    old_len = len(g.ndata)
    g.ndata.pop("h")
    assert len(g.ndata) == old_len - 1
    g.ndata["h"] = F.zeros((10, D))
    # set partial nodes
    u = F.tensor([1, 3, 5], g.idtype)
    g.nodes[u].data["h"] = F.ones((3, D))
    assert _pfc(g.ndata["h"]) == [
        0.0,
        1.0,
        0.0,
        1.0,
        0.0,
        1.0,
        0.0,
        0.0,
        0.0,
        0.0,
    ]
    # get partial nodes
    u = F.tensor([1, 2, 3], g.idtype)
    assert _pfc(g.nodes[u].data["h"]) == [1.0, 0.0, 1.0]

    """
    s, d, eid
    0, 1, 0
    1, 9, 1
    0, 2, 2
    2, 9, 3
    0, 3, 4
    3, 9, 5
    0, 4, 6
    4, 9, 7
    0, 5, 8
    5, 9, 9
    0, 6, 10
    6, 9, 11
    0, 7, 12
    7, 9, 13
    0, 8, 14
    8, 9, 15
    9, 0, 16
    """
    # set all edges
    g.edata["l"] = F.zeros((17, D))
    assert _pfc(g.edata["l"]) == [0.0] * 17
    # pop edges
    old_len = len(g.edata)
    g.edata.pop("l")