code: def _broadcast_shape(
      data, rank, world_size, num_parts, is_feat_data, feat_name
  ):
      """Auxiliary function to broadcast the shape of a feature data.
      This information is used to figure out the type-ids for the
      local features.
  
      Parameters:
      -----------
      data : numpy array
          which is the feature data read from the disk
      rank : integer
          which represents the id of the process in the process group
      world_size : integer
          represents the total no. of process in the process group
      num_parts : integer
          specifying the no. of partitions
      is_feat_data : bool
          flag used to seperate feature data and edge data
      feat_name : string
          name of the feature
  
      Returns:
      -------
      list of tuples :
          which represents the range of type-ids for the data array.
      """
      assert len(data.shape) in [
          1,
          2,
      ], f"Data is expected to be 1-D or 2-D but got {data.shape}."
      data_shape = list(data.shape)
  
      if len(data_shape) == 1:
          data_shape.append(1)
  
      if is_feat_data:
          data_shape.append(DATA_TYPE_ID[data.dtype])
  
      data_shape = torch.tensor(data_shape, dtype=torch.int64)
      data_shape_output = [
          torch.zeros_like(data_shape) for _ in range(world_size)
      ]
      dist.all_gather(data_shape_output, data_shape)
      logging.debug(
          f"[Rank: {rank} Received shapes from all ranks: {data_shape_output}"
      )
      shapes = [x.numpy() for x in data_shape_output if x[0] != 0]
      shapes = np.vstack(shapes)
  
      if is_feat_data:
          logging.debug(
              f"shapes: {shapes}, condition: {all(shapes[0,2] == s for s in shapes[:,2])}"
          )
          assert all(
              shapes[0, 2] == s for s in shapes[:, 2]
          ), f"dtypes for {feat_name} does not match on all ranks"
  
      # compute tids here.
      type_counts = list(shapes[:, 0])
      tid_start = np.cumsum([0] + type_counts[:-1])
      tid_end = np.cumsum(type_counts)
      tid_ranges = list(zip(tid_start, tid_end))
      logging.debug(f"starts -> {tid_start} ... end -> {tid_end}")
  
      return tid_ranges
  
file_path: dgl/tools/distpartitioning/dataset_utils.py
imports: import gc
  import logging
  import os
  import array_readwriter
  import constants
  import numpy
  import pyarrow
  import pyarrow.parquet
  import torch
  import torch.distributed
  from gloo_wrapper import alltoallv_cpu
  from utils import DATA_TYPE_ID, generate_read_list, get_gid_offsets, get_idranges, map_partid_rank, REV_DATA_TYPE_ID
label_desc: Use image node_id 1 for calling a global function with example usage: _broadcast_shape(data, rank, world_size, num_parts, is_feat_data, feat_name) and returns: tid_ranges