def test_features(idtype):
    """Test the features of batched DGLGraphs"""
    g1 = dgl.heterograph(
        {
            ("user", "follows", "user"): ([0, 1], [1, 2]),
            ("user", "plays", "game"): ([0, 1], [0, 0]),
        },
        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.nodes["game"].data["h1"] = F.tensor([[0.0]])
    g1.nodes["game"].data["h2"] = F.tensor([[1.0]])
    g1.edges["follows"].data["h1"] = F.tensor([[0.0], [1.0]])
    g1.edges["follows"].data["h2"] = F.tensor([[2.0], [3.0]])
    g1.edges["plays"].data["h1"] = F.tensor([[0.0], [1.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]])

    # test default setting
    bg = dgl.batch([g1, g2])
    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
        ),
    )
    assert F.allclose(
        bg.nodes["game"].data["h1"],
        F.cat(
            [g1.nodes["game"].data["h1"], g2.nodes["game"].data["h1"]], dim=0
        ),
    )
    assert F.allclose(
        bg.nodes["game"].data["h2"],
        F.cat(
            [g1.nodes["game"].data["h2"], g2.nodes["game"].data["h2"]], dim=0
        ),
    )
    assert F.allclose(
        bg.edges["follows"].data["h1"],
        F.cat(
            [g1.edges["follows"].data["h1"], g2.edges["follows"].data["h1"]],
            dim=0,
        ),
    )
    assert F.allclose(
        bg.edges["follows"].data["h2"],
        F.cat(
            [g1.edges["follows"].data["h2"], g2.edges["follows"].data["h2"]],
            dim=0,
        ),
    )
    assert F.allclose(
        bg.edges["plays"].data["h1"],
        F.cat(
            [g1.edges["plays"].data["h1"], g2.edges["plays"].data["h1"]], dim=0
        ),
    )

    # test specifying ndata/edata
    bg = dgl.batch([g1, g2], ndata=["h2"], edata=["h1"])
    assert F.allclose(
        bg.nodes["user"].data["h2"],
        F.cat(
            [g1.nodes["user"].data["h2"], g2.nodes["user"].data["h2"]], dim=0
        ),
    )
    assert F.allclose(
        bg.nodes["game"].data["h2"],
        F.cat(
            [g1.nodes["game"].data["h2"], g2.nodes["game"].data["h2"]], dim=0
        ),
    )
    assert F.allclose(
        bg.edges["follows"].data["h1"],
        F.cat(
            [g1.edges["follows"].data["h1"], g2.edges["follows"].data["h1"]],
            dim=0,
        ),
    )
    assert F.allclose(
        bg.edges["plays"].data["h1"],
        F.cat(
            [g1.edges["plays"].data["h1"], g2.edges["plays"].data["h1"]], dim=0
        ),
    )
    assert "h1" not in bg.nodes["user"].data
    assert "h1" not in bg.nodes["game"].data
    assert "h2" not in bg.edges["follows"].data

    # Test unbatching graphs
    g3, g4 = dgl.unbatch(bg)
    check_equivalence_between_heterographs(
        g1,
        g3,
        node_attrs={"user": ["h2"], "game": ["h2"]},
        edge_attrs={("user", "follows", "user"): ["h1"]},
    )
    check_equivalence_between_heterographs(
        g2,
        g4,
        node_attrs={"user": ["h2"], "game": ["h2"]},
        edge_attrs={("user", "follows", "user"): ["h1"]},
    )
Test the features of batched DGLGraphs
tests/python/common/test_batch-heterograph.py