def _verify_hetero_graph_node_edge_num(
    g,
    parts,
    store_inner_edge,
    debug_mode,
):
    """
    check list:
        make sure edge type are correct.
        make sure the number of nodes in each node type are correct.
        make sure the number of nodes in each node type are correct.
    """
    num_nodes = {ntype: 0 for ntype in g.ntypes}
    num_edges = {etype: 0 for etype in g.canonical_etypes}
    for part in parts:
        edata = (
            part.edge_attributes
            if isinstance(part, gb.FusedCSCSamplingGraph)
            else part.edata
        )
        if dgl.ETYPE in edata:
            # edata may not contain all edge types.
            assert len(g.canonical_etypes) >= len(F.unique(edata[dgl.ETYPE]))
        if debug_mode or isinstance(part, dgl.DGLGraph):
            for ntype in g.ntypes:
                ntype_id = g.get_ntype_id(ntype)
                inner_node_mask = _get_inner_node_mask(part, ntype_id)
                num_inner_nodes = F.sum(F.astype(inner_node_mask, F.int64), 0)
                num_nodes[ntype] += num_inner_nodes
        if store_inner_edge or isinstance(part, dgl.DGLGraph):
            for etype in g.canonical_etypes:
                etype_id = g.get_etype_id(etype)
                inner_edge_mask = _get_inner_edge_mask(part, etype_id)
                num_inner_edges = F.sum(F.astype(inner_edge_mask, F.int64), 0)
                num_edges[etype] += num_inner_edges

    # Verify the number of nodes are correct.
    if debug_mode or isinstance(part, dgl.DGLGraph):
        for ntype in g.ntypes:
            print(
                "node {}: {}, {}".format(
                    ntype, g.num_nodes(ntype), num_nodes[ntype]
                )
            )
            assert g.num_nodes(ntype) == num_nodes[ntype]
    # Verify the number of edges are correct.
    if store_inner_edge or isinstance(part, dgl.DGLGraph):
        for etype in g.canonical_etypes:
            print(
                "edge {}: {}, {}".format(
                    etype, g.num_edges(etype), num_edges[etype]
                )
            )
            assert g.num_edges(etype) == num_edges[etype]

    check list:
        make sure edge type are correct.
        make sure the number of nodes in each node type are correct.
        make sure the number of nodes in each node type are correct.
    
tests/distributed/test_partition.py