return a list of ntypes/etypes (#2742)

Co-authored-by: Da Zheng <zhengda1936@gmail.com>

python/dgl/distributed/partition.py

@@ -113,6 +113,7 @@ def load_partition(part_config, part_id):
     assert EID in graph.edata, "the partition graph should contain edge mapping to global edge ID"
 
     gpb, graph_name, ntypes, etypes = load_partition_book(part_config, part_id, graph)
+    ntypes_list, etypes_list = [], []
     for ntype in ntypes:
         ntype_id = ntypes[ntype]
         # graph.ndata[NID] are global homogeneous node IDs.
@@ -122,6 +123,7 @@ def load_partition(part_config, part_id):
         partids2 = gpb.nid2partid(per_type_nids, ntype)
         assert np.all(F.asnumpy(partids1 == part_id)), 'load a wrong partition'
         assert np.all(F.asnumpy(partids2 == part_id)), 'load a wrong partition'
+        ntypes_list.append(ntype)
     for etype in etypes:
         etype_id = etypes[etype]
         # graph.edata[EID] are global homogeneous edge IDs.
@@ -131,7 +133,8 @@ def load_partition(part_config, part_id):
         partids2 = gpb.eid2partid(per_type_eids, etype)
         assert np.all(F.asnumpy(partids1 == part_id)), 'load a wrong partition'
         assert np.all(F.asnumpy(partids2 == part_id)), 'load a wrong partition'
-    return graph, node_feats, edge_feats, gpb, graph_name, ntypes, etypes
+        etypes_list.append(etype)
+    return graph, node_feats, edge_feats, gpb, graph_name, ntypes_list, etypes_list
 
 def load_partition_book(part_config, part_id, graph=None):
     ''' Load a graph partition book from the partition config file.