Function test_node_frame

tests/python/common/test_heterograph-remove.py:172–180

(idtype)

@parametrize_idtype
def test_node_frame(idtype):
    g = create_graph(idtype, 10)
    data = np.random.rand(10, 3)
    new_data = data.take([0, 1, 2, 7, 8, 9], axis=0)
    g.ndata["h"] = F.tensor(data)

    # remove nodes
    g.remove_nodes(range(3, 7))
    assert F.allclose(g.ndata["h"], F.tensor(new_data))


@parametrize_idtype