Function test_uniform_random_walk

tests/python/common/sampling/test_sampling.py:150–266

(use_uva)

@pytest.mark.parametrize("use_uva", [True, False])
def test_uniform_random_walk(use_uva):
    if use_uva and F.ctx() == F.cpu():
        pytest.skip("UVA random walk requires a GPU.")
    g1 = dgl.heterograph({("user", "follow", "user"): ([0, 1, 2], [1, 2, 0])})
    g2 = dgl.heterograph(
        {("user", "follow", "user"): ([0, 1, 1, 2, 3], [1, 2, 3, 0, 0])}
    )
    g3 = dgl.heterograph(
        {
            ("user", "follow", "user"): ([0, 1, 2], [1, 2, 0]),
            ("user", "view", "item"): ([0, 1, 2], [0, 1, 2]),
            ("item", "viewed-by", "user"): ([0, 1, 2], [0, 1, 2]),
        }
    )
    g4 = dgl.heterograph(
        {
            ("user", "follow", "user"): ([0, 1, 1, 2, 3], [1, 2, 3, 0, 0]),
            ("user", "view", "item"): ([0, 0, 1, 2, 3, 3], [0, 1, 1, 2, 2, 1]),
            ("item", "viewed-by", "user"): (
                [0, 1, 1, 2, 2, 1],
                [0, 0, 1, 2, 3, 3],
            ),
        }
    )

    if use_uva:
        for g in (g1, g2, g3, g4):
            g.create_formats_()
            g.pin_memory_()
    elif F._default_context_str == "gpu":
        g1 = g1.to(F.ctx())
        g2 = g2.to(F.ctx())
        g3 = g3.to(F.ctx())
        g4 = g4.to(F.ctx())

    try:
        traces, eids, ntypes = dgl.sampling.random_walk(
            g1,
            F.tensor([0, 1, 2, 0, 1, 2], dtype=g1.idtype),
            length=4,
            return_eids=True,
        )
        check_random_walk(g1, ["follow"] * 4, traces, ntypes, trace_eids=eids)
        if F._default_context_str == "cpu":
            with pytest.raises(dgl.DGLError):
                dgl.sampling.random_walk(
                    g1,
                    F.tensor([0, 1, 2, 10], dtype=g1.idtype),
                    length=4,
                    return_eids=True,
                )
        traces, eids, ntypes = dgl.sampling.random_walk(
            g1,
            F.tensor([0, 1, 2, 0, 1, 2], dtype=g1.idtype),
            length=4,
            restart_prob=0.0,
            return_eids=True,
        )