code: def _set_trainer_ids(g, sim_g, node_parts):
      """Set the trainer IDs for each node and edge on the input graph.
  
      The trainer IDs will be stored as node data and edge data in the input graph.
  
      Parameters
      ----------
      g : DGLGraph
         The input graph for partitioning.
      sim_g : DGLGraph
          The homogeneous version of the input graph.
      node_parts : tensor
          The node partition ID for each node in `sim_g`.
      """
      if g.is_homogeneous:
          g.ndata["trainer_id"] = node_parts
          # An edge is assigned to a partition based on its destination node.
          g.edata["trainer_id"] = F.gather_row(node_parts, g.edges()[1])
      else:
          for ntype_id, ntype in enumerate(g.ntypes):
              type_idx = sim_g.ndata[NTYPE] == ntype_id
              orig_nid = F.boolean_mask(sim_g.ndata[NID], type_idx)
              trainer_id = F.zeros(
                  (len(orig_nid),), F.dtype(node_parts), F.cpu()
              )
              F.scatter_row_inplace(
                  trainer_id,
                  orig_nid,
                  F.boolean_mask(node_parts, type_idx),
              )
              g.nodes[ntype].data["trainer_id"] = trainer_id
          for c_etype in g.canonical_etypes:
              # An edge is assigned to a partition based on its destination node.
              _, _, dst_type = c_etype
              trainer_id = F.gather_row(
                  g.nodes[dst_type].data["trainer_id"],
                  g.edges(etype=c_etype)[1],
              )
              g.edges[c_etype].data["trainer_id"] = trainer_id
  
imports: import json
  import logging
  import os
  import time
  import numpy
  import torch
  from ..None import backend
  from ..base import DGLError, EID, ETYPE, NID, NTYPE
  from ..convert import to_homogeneous
  from ..data.utils import load_graphs, load_tensors, save_graphs, save_tensors
  from ..partition import get_peak_mem, metis_partition_assignment, partition_graph_with_halo
  from ..random import choice
  from ..transforms import sort_csc_by_tag, sort_csr_by_tag
  from .constants import DEFAULT_ETYPE, DEFAULT_NTYPE
  from .graph_partition_book import _etype_str_to_tuple, _etype_tuple_to_str, RangePartitionBook
label_desc: Use image node_id 12 for calling a global function with example usage: _set_trainer_ids(g, sim_g, node_parts) without return types