code: def _build_key_size_numel_dictionaries(keys, data):
      """Build the size on rank 0 and broadcast."""
      max_dim = _MAX_DATA_DIM
      sizes = [0 for _ in range(max_dim) for _ in keys]
  
      # Pack the sizes on rank zero.
      if get_tensor_model_parallel_rank() == 0:
          offset = 0
          for key in keys:
              assert (
                  data[key].dim() < max_dim
              ), "you should increase MAX_DATA_DIM"
              size = data[key].size()
              for i, s in enumerate(size):
                  sizes[i + offset] = s
              offset += max_dim
  
      # Move to GPU and broadcast.
      sizes_cuda = torch.cuda.LongTensor(sizes)
      torch.distributed.broadcast(
          sizes_cuda,
          get_tensor_model_parallel_src_rank(),
          group=get_tensor_model_parallel_group(),
      )
  
      # Move back to cpu and unpack.
      sizes_cpu = sizes_cuda.cpu()
      key_size = {}
      key_numel = {}
      total_numel = 0
      offset = 0
      for key in keys:
          i = 0
          size = []
          numel = 1
          while sizes_cpu[offset + i] > 0:
              this_size = sizes_cpu[offset + i]
              size.append(this_size)
              numel *= this_size
              i += 1
          key_size[key] = size
          key_numel[key] = numel
          total_numel += numel
          offset += max_dim
  
      return key_size, key_numel, total_numel
  
imports: import torch
  from apex.transformer.parallel_state import get_tensor_model_parallel_group
  from apex.transformer.parallel_state import get_tensor_model_parallel_rank
  from apex.transformer.parallel_state import get_tensor_model_parallel_src_rank
label_desc: Use image node_id 2 for calling a global function with example usage: _build_key_size_numel_dictionaries(keys, data) and returns: key_size,  key_numel,  total_numel