code: def reshuffle_graph(g, node_part=None):
      """Reshuffle node ids and edge IDs of a graph.
  
      This function reshuffles nodes and edges in a graph so that all nodes/edges of the same type
      have contiguous IDs. If a graph is partitioned and nodes are assigned to different partitions,
      all nodes/edges in a partition should
      get contiguous IDs; within a partition, all nodes/edges of the same type have contigous IDs.
  
      Parameters
      ----------
      g : DGLGraph
          The input graph.
      node_part : Tensor
          This is a vector whose length is the same as the number of nodes in the input graph.
          Each element indicates the partition ID the corresponding node is assigned to.
  
      Returns
      -------
      (DGLGraph, Tensor)
          The graph whose nodes and edges are reshuffled.
          The 1D tensor that indicates the partition IDs of the nodes in the reshuffled graph.
      """
      # In this case, we don't need to reshuffle node IDs and edge IDs.
      if node_part is None:
          g.ndata["orig_id"] = F.arange(0, g.num_nodes())
          g.edata["orig_id"] = F.arange(0, g.num_edges())
          return g, None
  
      start = time.time()
      if node_part is not None:
          node_part = utils.toindex(node_part)
          node_part = node_part.tousertensor()
      if NTYPE in g.ndata:
          is_hetero = len(F.unique(g.ndata[NTYPE])) > 1
      else:
          is_hetero = False
      if is_hetero:
          num_node_types = F.max(g.ndata[NTYPE], 0) + 1
          if node_part is not None:
              sorted_part, new2old_map = F.sort_1d(
                  node_part * num_node_types + g.ndata[NTYPE]
              )
          else:
              sorted_part, new2old_map = F.sort_1d(g.ndata[NTYPE])
          sorted_part = F.floor_div(sorted_part, num_node_types)
      elif node_part is not None:
          sorted_part, new2old_map = F.sort_1d(node_part)
      else:
          g.ndata["orig_id"] = g.ndata[NID]
          g.edata["orig_id"] = g.edata[EID]
          return g, None
  
      new_node_ids = np.zeros((g.num_nodes(),), dtype=np.int64)
      new_node_ids[F.asnumpy(new2old_map)] = np.arange(0, g.num_nodes())
      # If the input graph is homogneous, we only need to create an empty array, so that
      # _CAPI_DGLReassignEdges_Hetero knows how to handle it.
      etype = (
          g.edata[ETYPE]
          if ETYPE in g.edata
          else F.zeros((0), F.dtype(sorted_part), F.cpu())
      )
      g = reorder_nodes(g, new_node_ids)
      node_part = utils.toindex(sorted_part)
      # We reassign edges in in-CSR. In this way, after partitioning, we can ensure
      # that all edges in a partition are in the contiguous ID space.
      etype_idx = utils.toindex(etype)
      orig_eids = _CAPI_DGLReassignEdges_Hetero(
          g._graph,
          etype_idx.todgltensor(),
          node_part.todgltensor(),
          True,
      )
      orig_eids = utils.toindex(orig_eids)
      orig_eids = orig_eids.tousertensor()
      g.edata["orig_id"] = orig_eids
  
      print(
          "Reshuffle nodes and edges: {:.3f} seconds".format(
              time.time() - start
          )
      )
      return g, node_part.tousertensor()
  
imports: import os
  import re
  import time
  import numpy
  from .None import backend, utils
  from ._ffi.function import _init_api
  from .base import EID, ETYPE, NID, NTYPE
  from .heterograph import DGLGraph
  from .ndarray import NDArray
  from .subgraph import edge_subgraph
label_desc: Use image node_id 3 for calling a global function with example usage: reshuffle_graph(g, node_part) and returns: g,  node_part, g,  None, g,  None