def _verify_dgl_partition(graph, part_id, gpb, ntypes, etypes):
    """Verify the partition of a DGL graph."""
    assert (
        NID in graph.ndata
    ), "the partition graph should contain node mapping to global node ID"
    assert (
        EID in graph.edata
    ), "the partition graph should contain edge mapping to global edge ID"

    for ntype in ntypes:
        ntype_id = ntypes[ntype]
        # graph.ndata[NID] are global homogeneous node IDs.
        nids = F.boolean_mask(
            graph.ndata[NID], _get_inner_node_mask(graph, ntype_id)
        )
        partids1 = gpb.nid2partid(nids)
        _, per_type_nids = gpb.map_to_per_ntype(nids)
        partids2 = gpb.nid2partid(per_type_nids, ntype)
        assert np.all(F.asnumpy(partids1 == part_id)), (
            "Unexpected partition IDs are found in the loaded partition "
            "while querying via global homogeneous node IDs."
        )
        assert np.all(F.asnumpy(partids2 == part_id)), (
            "Unexpected partition IDs are found in the loaded partition "
            "while querying via type-wise node IDs."
        )
    for etype in etypes:
        etype_id = etypes[etype]
        # graph.edata[EID] are global homogeneous edge IDs.
        eids = F.boolean_mask(
            graph.edata[EID], _get_inner_edge_mask(graph, etype_id)
        )
        partids1 = gpb.eid2partid(eids)
        _, per_type_eids = gpb.map_to_per_etype(eids)
        partids2 = gpb.eid2partid(per_type_eids, etype)
        assert np.all(F.asnumpy(partids1 == part_id)), (
            "Unexpected partition IDs are found in the loaded partition "
            "while querying via global homogeneous edge IDs."
        )
        assert np.all(F.asnumpy(partids2 == part_id)), (
            "Unexpected partition IDs are found in the loaded partition "
            "while querying via type-wise edge IDs."
        )
Verify the partition of a DGL graph.
python/dgl/distributed/partition.py