Function test_batching_batched

tests/python/common/test_batch-heterograph.py:125–200

Test batching a DGLGraph and a batched DGLGraph.

(idtype)

@parametrize_idtype
def test_batching_batched(idtype):
    """Test batching a DGLGraph and a batched DGLGraph."""
    g1 = dgl.heterograph(
        {
            ("user", "follows", "user"): ([0, 1], [1, 2]),
            ("user", "plays", "game"): ([0, 1], [0, 0]),
        },
        idtype=idtype,
        device=F.ctx(),
    )
    g2 = dgl.heterograph(
        {
            ("user", "follows", "user"): ([0, 1], [1, 2]),
            ("user", "plays", "game"): ([0, 1], [0, 0]),
        },
        idtype=idtype,
        device=F.ctx(),
    )
    bg1 = dgl.batch([g1, g2])
    g3 = dgl.heterograph(
        {
            ("user", "follows", "user"): ([0], [1]),
            ("user", "plays", "game"): ([1], [0]),
        },
        idtype=idtype,
        device=F.ctx(),
    )
    bg2 = dgl.batch([bg1, g3])
    assert bg2.idtype == idtype
    assert bg2.device == F.ctx()
    assert bg2.ntypes == g3.ntypes
    assert bg2.etypes == g3.etypes
    assert bg2.canonical_etypes == g3.canonical_etypes
    assert bg2.batch_size == 3

    # Test number of nodes
    for ntype in bg2.ntypes:
        assert F.asnumpy(bg2.batch_num_nodes(ntype)).tolist() == [
            g1.num_nodes(ntype),
            g2.num_nodes(ntype),
            g3.num_nodes(ntype),
        ]
        assert bg2.num_nodes(ntype) == (
            g1.num_nodes(ntype) + g2.num_nodes(ntype) + g3.num_nodes(ntype)
        )

    # Test number of edges
    for etype in bg2.canonical_etypes:
        assert F.asnumpy(bg2.batch_num_edges(etype)).tolist() == [
            g1.num_edges(etype),
            g2.num_edges(etype),
            g3.num_edges(etype),
        ]
        assert bg2.num_edges(etype) == (
            g1.num_edges(etype) + g2.num_edges(etype) + g3.num_edges(etype)
        )

    # Test relabeled nodes