code: def etype2canonical_etype(part_config, config):
      num_parts = config[NUM_PARTS_KEY]
      edge_map = config[EDGE_MAP_KEY]
      etypes = list(edge_map.keys())
      # Get part id of each seed edge.
      partition_ids = []
      for _, bound in edge_map.items():
          for i in range(num_parts):
              if bound[i][1] > bound[i][0]:
                  partition_ids.append(i)
                  break
      partition_ids = torch.tensor(partition_ids)
  
      # Get starting index of each partition.
      shifts = []
      for i in range(num_parts):
          shifts.append(edge_map[etypes[0]][i][0])
      shifts = torch.tensor(shifts)
  
      canonical_etypes = {}
      part_ids = [
          part_id
          for part_id in range(num_parts)
          if part_id in partition_ids
      ]
      for part_id in part_ids:
          seed_etypes = [
              etypes[i]
              for i in range(len(etypes))
              if partition_ids[i] == part_id
          ]
          c_etype = _find_c_etypes_in_partition(
              part_id,
              seed_etypes,
              config[ETYPES_KEY],
              config[NTYPES_KEY],
              edge_map,
              shifts,
              part_config,
          )
          canonical_etypes.update(c_etype)
      return canonical_etypes
  
file_path: dgl/tools/change_etype_to_canonical_etype.py
imports: import argparse
  import json
  import logging
  import os
  import time
  import dgl
  import torch
  from dgl._ffi.base import DGLError
  from dgl.data.utils import load_graphs
  from dgl.utils import toindex
label_desc: Use image node_id 2 for calling a global function with example usage: etype2canonical_etype(part_config, config) and returns: canonical_etypes