Function test_empty_relation

tests/python/common/test_batch-heterograph.py:337–428

Test the features of batched DGLGraphs

(idtype)

)
@parametrize_idtype
def test_empty_relation(idtype):
    """Test the features of batched DGLGraphs"""
    g1 = dgl.heterograph(
        {
            ("user", "follows", "user"): ([0, 1], [1, 2]),
            ("user", "plays", "game"): ([], []),
        },
        idtype=idtype,
        device=F.ctx(),
    )
    g1.nodes["user"].data["h1"] = F.tensor([[0.0], [1.0], [2.0]])
    g1.nodes["user"].data["h2"] = F.tensor([[3.0], [4.0], [5.0]])
    g1.edges["follows"].data["h1"] = F.tensor([[0.0], [1.0]])
    g1.edges["follows"].data["h2"] = F.tensor([[2.0], [3.0]])

    g2 = dgl.heterograph(
        {
            ("user", "follows", "user"): ([0, 1], [1, 2]),
            ("user", "plays", "game"): ([0, 1], [0, 0]),
        },
        idtype=idtype,
        device=F.ctx(),
    )
    g2.nodes["user"].data["h1"] = F.tensor([[0.0], [1.0], [2.0]])
    g2.nodes["user"].data["h2"] = F.tensor([[3.0], [4.0], [5.0]])
    g2.nodes["game"].data["h1"] = F.tensor([[0.0]])
    g2.nodes["game"].data["h2"] = F.tensor([[1.0]])
    g2.edges["follows"].data["h1"] = F.tensor([[0.0], [1.0]])
    g2.edges["follows"].data["h2"] = F.tensor([[2.0], [3.0]])
    g2.edges["plays"].data["h1"] = F.tensor([[0.0], [1.0]])

    bg = dgl.batch([g1, g2])

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

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

    # Test features
    assert F.allclose(
        bg.nodes["user"].data["h1"],
        F.cat(
            [g1.nodes["user"].data["h1"], g2.nodes["user"].data["h1"]], dim=0
        ),
    )
    assert F.allclose(
        bg.nodes["user"].data["h2"],
        F.cat(
            [g1.nodes["user"].data["h2"], g2.nodes["user"].data["h2"]], dim=0