code: def allgather_sizes(
      send_data, world_size, num_parts, return_sizes=False
  ):
      """
      Perform all gather on list lengths, used to compute prefix sums
      to determine the offsets on each ranks. This is used to allocate
      global ids for edges/nodes on each ranks.
  
      Parameters
      ----------
      send_data : numpy array
          Data on which allgather is performed.
      world_size : integer
          No. of processes configured for execution
      num_parts : integer
          No. of output graph partitions
      return_sizes : bool
          Boolean flag to indicate whether to return raw sizes from each process
          or perform prefix sum on the raw sizes.
  
      Returns :
      ---------
          numpy array
              array with the prefix sum
      """
  
      # Assert on the world_size, num_parts
      assert (num_parts % world_size) == 0
  
      # compute the length of the local data
      send_length = len(send_data)
      out_tensor = torch.as_tensor(send_data, dtype=torch.int64)
      in_tensor = [
          torch.zeros(send_length, dtype=torch.int64)
          for _ in range(world_size)
      ]
  
      # all_gather message
      dist.all_gather(in_tensor, out_tensor)
  
      # Return on the raw sizes from each process
      if return_sizes:
          return torch.cat(in_tensor).numpy()
  
      # gather sizes in on array to return to the invoking function
      rank_sizes = np.zeros(num_parts + 1, dtype=np.int64)
      part_counts = torch.cat(in_tensor).numpy()
  
      count = rank_sizes[0]
      idx = 1
      for local_part_id in range(num_parts // world_size):
          for r in range(world_size):
              count += part_counts[
                  r * (num_parts // world_size) + local_part_id
              ]
              rank_sizes[idx] = count
              idx += 1
  
      return rank_sizes
  
imports: import numpy
  import torch
  import torch.distributed
label_desc: Use image node_id 1 for calling a global function with example usage: allgather_sizes(send_data, world_size, num_parts, return_sizes) and returns: rank_sizes, torch