Function _test

tests/python/common/test_heterograph-kernel.py:91–148

(red, partial)

def test_copy_src_reduce():
    def _test(red, partial):
        g = dgl.from_networkx(nx.erdos_renyi_graph(100, 0.1))
        # NOTE(zihao): add self-loop to avoid zero-degree nodes.
        # https://github.com/dmlc/dgl/issues/761
        g.add_edges(g.nodes(), g.nodes())
        g = g.to(F.ctx())
        hu, hv, he = generate_feature(g, "none", "none")
        if partial:
            nid = F.tensor(list(range(0, 100, 2)), g.idtype)

        g.ndata["u"] = F.attach_grad(F.clone(hu))
        g.ndata["v"] = F.attach_grad(F.clone(hv))
        g.edata["e"] = F.attach_grad(F.clone(he))

        with F.record_grad():
            if partial:
                g.pull(
                    nid,
                    fn.copy_u(u="u", out="m"),
                    builtin[red](msg="m", out="r1"),
                )
            else:
                g.update_all(
                    fn.copy_u(u="u", out="m"), builtin[red](msg="m", out="r1")
                )
            r1 = g.ndata["r1"]
            F.backward(F.reduce_sum(r1))
            n_grad1 = F.grad(g.ndata["u"])

        # reset grad
        g.ndata["u"] = F.attach_grad(F.clone(hu))
        g.ndata["v"] = F.attach_grad(F.clone(hv))
        g.edata["e"] = F.attach_grad(F.clone(he))

        with F.record_grad():
            if partial:
                g.pull(nid, udf_copy_src, udf_reduce[red])
            else:
                g.update_all(udf_copy_src, udf_reduce[red])
            r2 = g.ndata["r2"]
            F.backward(F.reduce_sum(r2))
            n_grad2 = F.grad(g.ndata["u"])

        def _print_error(a, b):
            print("ERROR: Test copy_src_{} partial: {}".format(red, partial))
            for i, (x, y) in enumerate(
                zip(F.asnumpy(a).flatten(), F.asnumpy(b).flatten())
            ):
                if not np.allclose(x, y):
                    print("@{} {} v.s. {}".format(i, x, y))

        if not F.allclose(r1, r2):
            _print_error(r1, r2)
        assert F.allclose(r1, r2)
        if not F.allclose(n_grad1, n_grad2):
            print("node grad")
            _print_error(n_grad1, n_grad2)
        assert F.allclose(n_grad1, n_grad2)