code: def check_hetero_partition(
      hg,
      part_method,
      num_parts=4,
      num_trainers_per_machine=1,
      load_feats=True,
      graph_formats=None,
  ):
      test_ntype = "n1"
      test_etype = ("n1", "r1", "n2")
      hg.nodes[test_ntype].data["labels"] = F.arange(
          0, hg.num_nodes(test_ntype)
      )
      hg.nodes[test_ntype].data["feats"] = F.tensor(
          np.random.randn(hg.num_nodes(test_ntype), 10), F.float32
      )
      hg.edges[test_etype].data["feats"] = F.tensor(
          np.random.randn(hg.num_edges(test_etype), 10), F.float32
      )
      hg.edges[test_etype].data["labels"] = F.arange(
          0, hg.num_edges(test_etype)
      )
      num_hops = 1
  
      orig_nids, orig_eids = partition_graph(
          hg,
          "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,
      )
      assert len(orig_nids) == len(hg.ntypes)
      assert len(orig_eids) == len(hg.canonical_etypes)
      for ntype in hg.ntypes:
          assert len(orig_nids[ntype]) == hg.num_nodes(ntype)
      for etype in hg.canonical_etypes:
          assert len(orig_eids[etype]) == hg.num_edges(etype)
      parts = []
      shuffled_labels = []
      shuffled_elabels = []
      for i in range(num_parts):
          (
              part_g,
              node_feats,
              edge_feats,
              gpb,
              _,
              ntypes,
              etypes,
          ) = 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 hg.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 hg.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)
          # Verify the mapping between the reshuffled IDs and the original IDs.
          # These are partition-local IDs.
          part_src_ids, part_dst_ids = part_g.edges()
          # These are reshuffled global homogeneous IDs.
          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]
          # These are reshuffled per-type IDs.
          src_ntype_ids, part_src_ids = gpb.map_to_per_ntype(
              part_src_ids
          )
          dst_ntype_ids, part_dst_ids = gpb.map_to_per_ntype(
              part_dst_ids
          )
          etype_ids, part_eids = gpb.map_to_per_etype(part_eids)
          # These are original per-type IDs.
          for etype_id, etype in enumerate(hg.canonical_etypes):
              part_src_ids1 = F.boolean_mask(
                  part_src_ids, etype_ids == etype_id
              )
              src_ntype_ids1 = F.boolean_mask(
                  src_ntype_ids, etype_ids == etype_id
              )
              part_dst_ids1 = F.boolean_mask(
                  part_dst_ids, etype_ids == etype_id
              )
              dst_ntype_ids1 = F.boolean_mask(
                  dst_ntype_ids, etype_ids == etype_id
              )
              part_eids1 = F.boolean_mask(
                  part_eids, etype_ids == etype_id
              )
              assert np.all(
                  F.asnumpy(src_ntype_ids1 == src_ntype_ids1[0])
              )
              assert np.all(
                  F.asnumpy(dst_ntype_ids1 == dst_ntype_ids1[0])
              )
              src_ntype = hg.ntypes[F.as_scalar(src_ntype_ids1[0])]
              dst_ntype = hg.ntypes[F.as_scalar(dst_ntype_ids1[0])]
              orig_src_ids1 = F.gather_row(
                  orig_nids[src_ntype], part_src_ids1
              )
              orig_dst_ids1 = F.gather_row(
                  orig_nids[dst_ntype], part_dst_ids1
              )
              orig_eids1 = F.gather_row(orig_eids[etype], part_eids1)
              orig_eids2 = hg.edge_ids(
                  orig_src_ids1, orig_dst_ids1, etype=etype
              )
              assert len(orig_eids1) == len(orig_eids2)
              assert np.all(
                  F.asnumpy(orig_eids1) == F.asnumpy(orig_eids2)
              )
          parts.append(part_g)
          verify_graph_feats(
              hg,
              gpb,
              part_g,
              node_feats,
              edge_feats,
              orig_nids,
              orig_eids,
          )
  
          shuffled_labels.append(node_feats[test_ntype + "/labels"])
          shuffled_elabels.append(
              edge_feats[_etype_tuple_to_str(test_etype) + "/labels"]
          )
      verify_hetero_graph(hg, parts)
  
      shuffled_labels = F.asnumpy(F.cat(shuffled_labels, 0))
      shuffled_elabels = F.asnumpy(F.cat(shuffled_elabels, 0))
      orig_labels = np.zeros(
          shuffled_labels.shape, dtype=shuffled_labels.dtype
      )
      orig_elabels = np.zeros(
          shuffled_elabels.shape, dtype=shuffled_elabels.dtype
      )
      orig_labels[F.asnumpy(orig_nids[test_ntype])] = shuffled_labels
      orig_elabels[F.asnumpy(orig_eids[test_etype])] = shuffled_elabels
      assert np.all(
          orig_labels == F.asnumpy(hg.nodes[test_ntype].data["labels"])
      )
      assert np.all(
          orig_elabels == F.asnumpy(hg.edges[test_etype].data["labels"])
      )
  
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 7 for calling a global function with example usage: check_hetero_partition(hg, part_method, num_parts, num_trainers_per_machine, load_feats, graph_formats) without return types