code: def check_metis_partition(g, extra_hops):
      subgs = dgl.transforms.metis_partition(
          g, 4, extra_cached_hops=extra_hops
      )
      num_inner_nodes = 0
      num_inner_edges = 0
      if subgs is not None:
          for part_id, subg in subgs.items():
              lnode_ids = np.nonzero(
                  F.asnumpy(subg.ndata["inner_node"])
              )[0]
              ledge_ids = np.nonzero(
                  F.asnumpy(subg.edata["inner_edge"])
              )[0]
              num_inner_nodes += len(lnode_ids)
              num_inner_edges += len(ledge_ids)
              assert np.sum(
                  F.asnumpy(subg.ndata["part_id"]) == part_id
              ) == len(lnode_ids)
          assert num_inner_nodes == g.num_nodes()
          print(g.num_edges() - num_inner_edges)
  
      if extra_hops == 0:
          return
  
      # partitions with node reshuffling
      subgs = dgl.transforms.metis_partition(
          g, 4, extra_cached_hops=extra_hops, reshuffle=True
      )
      num_inner_nodes = 0
      num_inner_edges = 0
      edge_cnts = np.zeros((g.num_edges(),))
      if subgs is not None:
          for part_id, subg in subgs.items():
              lnode_ids = np.nonzero(
                  F.asnumpy(subg.ndata["inner_node"])
              )[0]
              ledge_ids = np.nonzero(
                  F.asnumpy(subg.edata["inner_edge"])
              )[0]
              num_inner_nodes += len(lnode_ids)
              num_inner_edges += len(ledge_ids)
              assert np.sum(
                  F.asnumpy(subg.ndata["part_id"]) == part_id
              ) == len(lnode_ids)
              nids = F.asnumpy(subg.ndata[dgl.NID])
  
              # ensure the local node Ids are contiguous.
              parent_ids = F.asnumpy(subg.ndata[dgl.NID])
              parent_ids = parent_ids[: len(lnode_ids)]
              assert np.all(
                  parent_ids
                  == np.arange(parent_ids[0], parent_ids[-1] + 1)
              )
  
              # count the local edges.
              parent_ids = F.asnumpy(subg.edata[dgl.EID])[ledge_ids]
              edge_cnts[parent_ids] += 1
  
              orig_ids = subg.ndata["orig_id"]
              inner_node = F.asnumpy(subg.ndata["inner_node"])
              for nid in range(subg.num_nodes()):
                  neighs = subg.predecessors(nid)
                  old_neighs1 = F.gather_row(orig_ids, neighs)
                  old_nid = F.asnumpy(orig_ids[nid])
                  old_neighs2 = g.predecessors(old_nid)
                  # If this is an inner node, it should have the full neighborhood.
                  if inner_node[nid]:
                      assert np.all(
                          np.sort(F.asnumpy(old_neighs1))
                          == np.sort(F.asnumpy(old_neighs2))
                      )
          # Normally, local edges are only counted once.
          assert np.all(edge_cnts == 1)
  
          assert num_inner_nodes == g.num_nodes()
          print(g.num_edges() - num_inner_edges)
  
imports: import math
  import os
  import unittest
  import backend
  import dgl
  import dgl.function
  import dgl.partition
  import networkx
  import numpy
  import pytest
  from scipy import sparse
  from utils import parametrize_idtype
  from utils.graph_cases import get_cases
label_desc: Use image node_id 20 for calling a global function with example usage: check_metis_partition(g, extra_hops) without return types