code: def check_partition(
      g,
      part_method,
      num_parts=4,
      num_trainers_per_machine=1,
      load_feats=True,
      graph_formats=None,
  ):
      g.ndata["labels"] = F.arange(0, g.num_nodes())
      g.ndata["feats"] = F.tensor(
          np.random.randn(g.num_nodes(), 10), F.float32
      )
      g.edata["feats"] = F.tensor(
          np.random.randn(g.num_edges(), 10), F.float32
      )
      g.update_all(fn.copy_u("feats", "msg"), fn.sum("msg", "h"))
      g.update_all(fn.copy_e("feats", "msg"), fn.sum("msg", "eh"))
      num_hops = 2
  
      orig_nids, orig_eids = partition_graph(
          g,
          "test",
          num_parts,
          "/tmp/partition",
          num_hops=num_hops,
          part_method=part_method,
          return_mapping=True,
          num_trainers_per_machine=num_trainers_per_machine,
          graph_formats=graph_formats,
      )
      part_sizes = []
      shuffled_labels = []
      shuffled_edata = []
      for i in range(num_parts):
          part_g, node_feats, edge_feats, gpb, _, _, _ = load_partition(
              "/tmp/partition/test.json", i, load_feats=load_feats
          )
          _verify_partition_data_types(part_g)
          _verify_partition_formats(part_g, graph_formats)
          if not load_feats:
              assert not node_feats
              assert not edge_feats
              node_feats, edge_feats = load_partition_feats(
                  "/tmp/partition/test.json", i
              )
          if num_trainers_per_machine > 1:
              for ntype in g.ntypes:
                  name = ntype + "/trainer_id"
                  assert name in node_feats
                  part_ids = F.floor_div(
                      node_feats[name], num_trainers_per_machine
                  )
                  assert np.all(F.asnumpy(part_ids) == i)
  
              for etype in g.canonical_etypes:
                  name = _etype_tuple_to_str(etype) + "/trainer_id"
                  assert name in edge_feats
                  part_ids = F.floor_div(
                      edge_feats[name], num_trainers_per_machine
                  )
                  assert np.all(F.asnumpy(part_ids) == i)
  
          # Check the metadata
          assert gpb._num_nodes() == g.num_nodes()
          assert gpb._num_edges() == g.num_edges()
  
          assert gpb.num_partitions() == num_parts
          gpb_meta = gpb.metadata()
          assert len(gpb_meta) == num_parts
          assert len(gpb.partid2nids(i)) == gpb_meta[i]["num_nodes"]
          assert len(gpb.partid2eids(i)) == gpb_meta[i]["num_edges"]
          part_sizes.append(
              (gpb_meta[i]["num_nodes"], gpb_meta[i]["num_edges"])
          )
  
          nid = F.boolean_mask(
              part_g.ndata[dgl.NID], part_g.ndata["inner_node"]
          )
          local_nid = gpb.nid2localnid(nid, i)
          assert F.dtype(local_nid) in (F.int64, F.int32)
          assert np.all(
              F.asnumpy(local_nid) == np.arange(0, len(local_nid))
          )
          eid = F.boolean_mask(
              part_g.edata[dgl.EID], part_g.edata["inner_edge"]
          )
          local_eid = gpb.eid2localeid(eid, i)
          assert F.dtype(local_eid) in (F.int64, F.int32)
          assert np.all(
              F.asnumpy(local_eid) == np.arange(0, len(local_eid))
          )
  
          # Check the node map.
          local_nodes = F.boolean_mask(
              part_g.ndata[dgl.NID], part_g.ndata["inner_node"]
          )
          llocal_nodes = F.nonzero_1d(part_g.ndata["inner_node"])
          local_nodes1 = gpb.partid2nids(i)
          assert F.dtype(local_nodes1) in (F.int32, F.int64)
          assert np.all(
              np.sort(F.asnumpy(local_nodes))
              == np.sort(F.asnumpy(local_nodes1))
          )
          assert np.all(
              F.asnumpy(llocal_nodes) == np.arange(len(llocal_nodes))
          )
  
          # Check the edge map.
          local_edges = F.boolean_mask(
              part_g.edata[dgl.EID], part_g.edata["inner_edge"]
          )
          llocal_edges = F.nonzero_1d(part_g.edata["inner_edge"])
          local_edges1 = gpb.partid2eids(i)
          assert F.dtype(local_edges1) in (F.int32, F.int64)
          assert np.all(
              np.sort(F.asnumpy(local_edges))
              == np.sort(F.asnumpy(local_edges1))
          )
          assert np.all(
              F.asnumpy(llocal_edges) == np.arange(len(llocal_edges))
          )
  
          # Verify the mapping between the reshuffled IDs and the original IDs.
          part_src_ids, part_dst_ids = part_g.edges()
          part_src_ids = F.gather_row(
              part_g.ndata[dgl.NID], part_src_ids
          )
          part_dst_ids = F.gather_row(
              part_g.ndata[dgl.NID], part_dst_ids
          )
          part_eids = part_g.edata[dgl.EID]
          orig_src_ids = F.gather_row(orig_nids, part_src_ids)
          orig_dst_ids = F.gather_row(orig_nids, part_dst_ids)
          orig_eids1 = F.gather_row(orig_eids, part_eids)
          orig_eids2 = g.edge_ids(orig_src_ids, orig_dst_ids)
          assert F.shape(orig_eids1)[0] == F.shape(orig_eids2)[0]
          assert np.all(F.asnumpy(orig_eids1) == F.asnumpy(orig_eids2))
  
          local_orig_nids = orig_nids[part_g.ndata[dgl.NID]]
          local_orig_eids = orig_eids[part_g.edata[dgl.EID]]
          part_g.ndata["feats"] = F.gather_row(
              g.ndata["feats"], local_orig_nids
          )
          part_g.edata["feats"] = F.gather_row(
              g.edata["feats"], local_orig_eids
          )
          local_nodes = orig_nids[local_nodes]
          local_edges = orig_eids[local_edges]
  
          part_g.update_all(
              fn.copy_u("feats", "msg"), fn.sum("msg", "h")
          )
          part_g.update_all(
              fn.copy_e("feats", "msg"), fn.sum("msg", "eh")
          )
          assert F.allclose(
              F.gather_row(g.ndata["h"], local_nodes),
              F.gather_row(part_g.ndata["h"], llocal_nodes),
          )
          assert F.allclose(
              F.gather_row(g.ndata["eh"], local_nodes),
              F.gather_row(part_g.ndata["eh"], llocal_nodes),
          )
  
          for name in ["labels", "feats"]:
              assert "_N/" + name in node_feats
              assert node_feats["_N/" + name].shape[0] == len(
                  local_nodes
              )
              true_feats = F.gather_row(g.ndata[name], local_nodes)
              ndata = F.gather_row(node_feats["_N/" + name], local_nid)
              assert np.all(F.asnumpy(true_feats) == F.asnumpy(ndata))
          for name in ["feats"]:
              efeat_name = (
                  _etype_tuple_to_str(DEFAULT_ETYPE) + "/" + name
              )
              assert efeat_name in edge_feats
              assert edge_feats[efeat_name].shape[0] == len(local_edges)
              true_feats = F.gather_row(g.edata[name], local_edges)
              edata = F.gather_row(edge_feats[efeat_name], local_eid)
              assert np.all(F.asnumpy(true_feats) == F.asnumpy(edata))
  
          # This only works if node/edge IDs are shuffled.
          shuffled_labels.append(node_feats["_N/labels"])
          shuffled_edata.append(edge_feats["_N:_E:_N/feats"])
  
      # Verify that we can reconstruct node/edge data for original IDs.
      shuffled_labels = F.asnumpy(F.cat(shuffled_labels, 0))
      shuffled_edata = F.asnumpy(F.cat(shuffled_edata, 0))
      orig_labels = np.zeros(
          shuffled_labels.shape, dtype=shuffled_labels.dtype
      )
      orig_edata = np.zeros(
          shuffled_edata.shape, dtype=shuffled_edata.dtype
      )
      orig_labels[F.asnumpy(orig_nids)] = shuffled_labels
      orig_edata[F.asnumpy(orig_eids)] = shuffled_edata
      assert np.all(orig_labels == F.asnumpy(g.ndata["labels"]))
      assert np.all(orig_edata == F.asnumpy(g.edata["feats"]))
  
      node_map = []
      edge_map = []
      for i, (num_nodes, num_edges) in enumerate(part_sizes):
          node_map.append(np.ones(num_nodes) * i)
          edge_map.append(np.ones(num_edges) * i)
      node_map = np.concatenate(node_map)
      edge_map = np.concatenate(edge_map)
      nid2pid = gpb.nid2partid(F.arange(0, len(node_map)))
      assert F.dtype(nid2pid) in (F.int32, F.int64)
      assert np.all(F.asnumpy(nid2pid) == node_map)
      eid2pid = gpb.eid2partid(F.arange(0, len(edge_map)))
      assert F.dtype(eid2pid) in (F.int32, F.int64)
      assert np.all(F.asnumpy(eid2pid) == edge_map)
  
imports: import json
  import os
  import tempfile
  import backend
  import dgl
  import numpy
  import pytest
  import torch
  from dgl import function
  from dgl.distributed import convert_dgl_partition_to_csc_sampling_graph, load_partition, load_partition_book, load_partition_feats, partition_graph
  from dgl.distributed.graph_partition_book import _etype_str_to_tuple, _etype_tuple_to_str, DEFAULT_ETYPE, DEFAULT_NTYPE, EdgePartitionPolicy, HeteroDataName, NodePartitionPolicy, RangePartitionBook
  from dgl.distributed.partition import _get_inner_edge_mask, _get_inner_node_mask, RESERVED_FIELD_DTYPE
  from scipy import sparse
  from utils import reset_envs
label_desc: Use image node_id 8 for calling a global function with example usage: check_partition(g, part_method, num_parts, num_trainers_per_machine, load_feats, graph_formats) without return types