Function test_copy_edge_reduce

tests/python/common/test_heterograph-kernel.py:158–223

()

def test_copy_edge_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.
        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_e(e="e", out="m"),
                    builtin[red](msg="m", out="r1"),
                )
            else:
                g.update_all(
                    fn.copy_e(e="e", out="m"), builtin[red](msg="m", out="r1")
                )
            r1 = g.ndata["r1"]
            F.backward(F.reduce_sum(r1))
            e_grad1 = F.grad(g.edata["e"])

        # 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_edge, udf_reduce[red])
            else:
                g.update_all(udf_copy_edge, udf_reduce[red])
            r2 = g.ndata["r2"]
            F.backward(F.reduce_sum(r2))
            e_grad2 = F.grad(g.edata["e"])

        def _print_error(a, b):
            print("ERROR: Test copy_edge_{} partial: {}".format(red, partial))
            return
            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(e_grad1, e_grad2):
            print("edge gradient")
            _print_error(e_grad1, e_grad2)