Function test_to_device

tests/python/common/test_heterograph.py:1135–1173

(idtype)

)
@parametrize_idtype
def test_to_device(idtype):
    # TODO: rewrite this test case to accept different graphs so we
    #  can test reverse graph and batched graph
    g = create_test_heterograph(idtype)
    g.nodes["user"].data["h"] = F.ones((3, 5))
    g.nodes["game"].data["i"] = F.ones((2, 5))
    g.edges["plays"].data["e"] = F.ones((4, 4))
    assert g.device == F.ctx()
    g = g.to(F.cpu())
    assert g.device == F.cpu()
    assert F.context(g.nodes["user"].data["h"]) == F.cpu()
    assert F.context(g.nodes["game"].data["i"]) == F.cpu()
    assert F.context(g.edges["plays"].data["e"]) == F.cpu()
    for ntype in g.ntypes:
        assert F.context(g.batch_num_nodes(ntype)) == F.cpu()
    for etype in g.canonical_etypes:
        assert F.context(g.batch_num_edges(etype)) == F.cpu()

    if F.is_cuda_available():
        g1 = g.to(F.cuda())
        assert g1.device == F.cuda()
        assert F.context(g1.nodes["user"].data["h"]) == F.cuda()
        assert F.context(g1.nodes["game"].data["i"]) == F.cuda()
        assert F.context(g1.edges["plays"].data["e"]) == F.cuda()
        for ntype in g1.ntypes:
            assert F.context(g1.batch_num_nodes(ntype)) == F.cuda()
        for etype in g1.canonical_etypes:
            assert F.context(g1.batch_num_edges(etype)) == F.cuda()
        assert F.context(g.nodes["user"].data["h"]) == F.cpu()
        assert F.context(g.nodes["game"].data["i"]) == F.cpu()
        assert F.context(g.edges["plays"].data["e"]) == F.cpu()
        for ntype in g.ntypes:
            assert F.context(g.batch_num_nodes(ntype)) == F.cpu()
        for etype in g.canonical_etypes:
            assert F.context(g.batch_num_edges(etype)) == F.cpu()
        with pytest.raises(DGLError):
            g1.nodes["user"].data["h"] = F.copy_to(F.ones((3, 5)), F.cpu())
        with pytest.raises(DGLError):
            g1.edges["plays"].data["e"] = F.copy_to(F.ones((4, 4)), F.cpu())


@unittest.skipIf(