Function test_pinsage_sampling

tests/python/common/sampling/test_sampling.py:316–398

(use_uva)

@pytest.mark.parametrize("use_uva", [True, False])
def test_pinsage_sampling(use_uva):
    if use_uva and F.ctx() == F.cpu():
        pytest.skip("UVA sampling requires a GPU.")

    def _test_sampler(g, sampler, ntype):
        seeds = F.copy_to(F.tensor([0, 2], dtype=g.idtype), F.ctx())
        neighbor_g = sampler(seeds)
        assert neighbor_g.ntypes == [ntype]
        u, v = neighbor_g.all_edges(form="uv", order="eid")
        uv = list(zip(F.asnumpy(u).tolist(), F.asnumpy(v).tolist()))
        assert (1, 0) in uv or (0, 0) in uv
        assert (2, 2) in uv or (3, 2) in uv

    g = dgl.heterograph(
        {
            ("item", "bought-by", "user"): (
                [0, 0, 1, 1, 2, 2, 3, 3],
                [0, 1, 0, 1, 2, 3, 2, 3],
            ),
            ("user", "bought", "item"): (
                [0, 1, 0, 1, 2, 3, 2, 3],
                [0, 0, 1, 1, 2, 2, 3, 3],
            ),
        }
    )
    if use_uva:
        g.create_formats_()
        g.pin_memory_()
    elif F._default_context_str == "gpu":
        g = g.to(F.ctx())
    try:
        sampler = dgl.sampling.PinSAGESampler(g, "item", "user", 4, 0.5, 3, 2)
        _test_sampler(g, sampler, "item")
        sampler = dgl.sampling.RandomWalkNeighborSampler(
            g, 4, 0.5, 3, 2, ["bought-by", "bought"]
        )
        _test_sampler(g, sampler, "item")
        sampler = dgl.sampling.RandomWalkNeighborSampler(
            g,
            4,
            0.5,
            3,
            2,
            [("item", "bought-by", "user"), ("user", "bought", "item")],
        )
        _test_sampler(g, sampler, "item")
    finally:
        if g.is_pinned():
            g.unpin_memory_()

    g = dgl.graph(([0, 0, 1, 1, 2, 2, 3, 3], [0, 1, 0, 1, 2, 3, 2, 3]))
    if use_uva:
        g.create_formats_()
        g.pin_memory_()
    elif F._default_context_str == "gpu":
        g = g.to(F.ctx())
    try:
        sampler = dgl.sampling.RandomWalkNeighborSampler(g, 4, 0.5, 3, 2)