Function test_local_var

tests/python/common/function/test_basics.py:521–590

(idtype)

@parametrize_idtype
def test_local_var(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):
        g = g.local_var()
        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):
        g = g.local_var()
        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):
        g = g.local_var()
        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):
        g = g.local_var()
        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 initializer1
    g = dgl.graph(([0, 1], [1, 1]), idtype=idtype, device=F.ctx())
    g.set_n_initializer(dgl.init.zero_initializer)

    def foo(g):
        g = g.local_var()
        g.nodes[0].data["h"] = F.ones((1, 1))
        assert F.allclose(g.ndata["h"], F.tensor([[1.0], [0.0]]))

    foo(g)

    # test initializer2
    def foo_e_initializer(shape, dtype, ctx, id_range):