code: def assign_shuffle_global_nids_edges(
      rank, world_size, num_parts, edge_data
  ):
      """
      Utility function to assign shuffle_global_eids to edges
      edge_data gets converted from [global_src_nid, global_dst_nid, global_type_eid, etype]
      to [shuffle_global_src_nid, shuffle_global_dst_nid, global_src_nid, global_dst_nid, global_type_eid, etype]
  
      Parameters:
      -----------
      rank : integer
          rank of the current process
      world_size : integer
          total count of processes in execution
      num_parts : integer
          total number of output graph partitions
      edge_data : numpy ndarray
          edge data as read from xxx_edges.txt file
  
      Returns:
      --------
      integer
          shuffle_global_eid_start, which indicates the starting value from which shuffle_global-ids are assigned to edges
          on this rank
      """
      # get prefix sum of edge counts per rank to locate the starting point
      # from which global-ids to edges are assigned in the current rank
      local_row_counts = []
      for local_part_id in range(num_parts // world_size):
          local_row_counts.append(
              edge_data[
                  constants.GLOBAL_SRC_ID + "/" + str(local_part_id)
              ].shape[0]
          )
  
      shuffle_global_eid_offset = []
      prefix_sum_edges = allgather_sizes(
          local_row_counts, world_size, num_parts
      )
      for local_part_id in range(num_parts // world_size):
          shuffle_global_eid_start = prefix_sum_edges[
              rank + (local_part_id * world_size)
          ]
          shuffle_global_eid_end = prefix_sum_edges[
              rank + 1 + (local_part_id * world_size)
          ]
          shuffle_global_eids = np.arange(
              shuffle_global_eid_start,
              shuffle_global_eid_end,
              dtype=np.int64,
          )
          edge_data[
              constants.SHUFFLE_GLOBAL_EID + "/" + str(local_part_id)
          ] = shuffle_global_eids
          shuffle_global_eid_offset.append(shuffle_global_eid_start)
  
      return shuffle_global_eid_offset
  
file_path: dgl/tools/distpartitioning/globalids.py
imports: import itertools
  import operator
  import constants
  import numpy
  import torch
  from dist_lookup import DistLookupService
  from gloo_wrapper import allgather_sizes, alltoallv_cpu
  from utils import memory_snapshot
label_desc: Use image node_id 4 for calling a global function with example usage: assign_shuffle_global_nids_edges(rank, world_size, num_parts, edge_data) and returns: shuffle_global_eid_offset