Function _test2

tests/python/common/sampling/test_sampling.py:615–657

(p, replace)

    _test1(prob, False)  # w/o replacement, uniform

    def _test2(p, replace):  # fanout > #neighbors
        subg = sample_neighbors_fusing_mode[fused](
            g, [0, 2], -1, prob=p, replace=replace
        )
        if not fused:
            assert subg.num_nodes() == g.num_nodes()
        u, v = subg.edges()
        if fused:
            u, v = subg.srcdata[dgl.NID][u], subg.dstdata[dgl.NID][v]
        u_ans, v_ans, e_ans = g.in_edges([0, 2], form="all")
        if p is not None:
            emask = F.gather_row(g.edata[p], e_ans)
            if p == "prob":
                emask = emask != 0
            u_ans = F.boolean_mask(u_ans, emask)
            v_ans = F.boolean_mask(v_ans, emask)
        uv = set(zip(F.asnumpy(u), F.asnumpy(v)))
        uv_ans = set(zip(F.asnumpy(u_ans), F.asnumpy(v_ans)))
        assert uv == uv_ans

        for i in range(10):
            subg = sample_neighbors_fusing_mode[fused](
                g, [0, 2], 2, prob=p, replace=replace
            )
            if not fused:
                assert subg.num_nodes() == g.num_nodes()
            num_edges = 4 if replace else 3
            assert subg.num_edges() == num_edges
            u, v = subg.edges()
            if fused:
                u, v = subg.srcdata[dgl.NID][u], subg.dstdata[dgl.NID][v]
            assert set(F.asnumpy(F.unique(v))) == {0, 2}
            assert F.array_equal(
                F.astype(g.has_edges_between(u, v), F.int64),
                F.ones((num_edges,), dtype=F.int64),
            )
            assert F.array_equal(g.edge_ids(u, v), subg.edata[dgl.EID])
            edge_set = set(zip(list(F.asnumpy(u)), list(F.asnumpy(v))))
            if not replace:
                # check no duplication
                assert len(edge_set) == num_edges
            if p is not None:
                assert not (3, 0) in edge_set

    _test2(prob, True)  # w/ replacement, uniform
    _test2(prob, False)  # w/o replacement, uniform