Function test_local_scope

tests/python/common/function/test_basics.py:594–691

(idtype)

@parametrize_idtype
def test_local_scope(idtype):
    g = dgl.graph(([0, 1, 2, 3], [1, 2, 3, 4]), idtype=idtype, device=F.ctx())
    g.ndata["h"] = F.zeros((g.num_nodes(), 3))
    g.edata["w"] = F.zeros((g.num_edges(), 4))

    # test override
    def foo(g):
        with g.local_scope():
            g.ndata["h"] = F.ones((g.num_nodes(), 3))
            g.edata["w"] = F.ones((g.num_edges(), 4))

    foo(g)
    assert F.allclose(g.ndata["h"], F.zeros((g.num_nodes(), 3)))
    assert F.allclose(g.edata["w"], F.zeros((g.num_edges(), 4)))

    # test out-place update
    def foo(g):
        with g.local_scope():
            g.nodes[[2, 3]].data["h"] = F.ones((2, 3))
            g.edges[[2, 3]].data["w"] = F.ones((2, 4))

    foo(g)
    assert F.allclose(g.ndata["h"], F.zeros((g.num_nodes(), 3)))
    assert F.allclose(g.edata["w"], F.zeros((g.num_edges(), 4)))

    # test out-place update 2
    def foo(g):
        with g.local_scope():
            g.apply_nodes(lambda nodes: {"h": nodes.data["h"] + 10}, [2, 3])
            g.apply_edges(lambda edges: {"w": edges.data["w"] + 10}, [2, 3])

    foo(g)
    assert F.allclose(g.ndata["h"], F.zeros((g.num_nodes(), 3)))
    assert F.allclose(g.edata["w"], F.zeros((g.num_edges(), 4)))

    # test auto-pop
    def foo(g):
        with g.local_scope():
            g.ndata["hh"] = F.ones((g.num_nodes(), 3))
            g.edata["ww"] = F.ones((g.num_edges(), 4))

    foo(g)
    assert "hh" not in g.ndata
    assert "ww" not in g.edata

    # test nested scope
    def foo(g):
        with g.local_scope():
            g.ndata["hh"] = F.ones((g.num_nodes(), 3))
            g.edata["ww"] = F.ones((g.num_edges(), 4))
            with g.local_scope():
                g.ndata["hhh"] = F.ones((g.num_nodes(), 3))
                g.edata["www"] = F.ones((g.num_edges(), 4))
            assert "hhh" not in g.ndata
            assert "www" not in g.edata

    foo(g)
    assert "hh" not in g.ndata