Function test_edge_frame

tests/python/common/test_heterograph-remove.py:184–193

(idtype)

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

    # remove edges
    g.remove_edges(range(3, 7))
    assert F.allclose(g.edata["h"], F.tensor(new_data))


@parametrize_idtype