code: def get_shard_size(total_size, mp_size, rank=None):
      global num_kv_heads
      # When we have num_kv_heads defined, uneven division is possible, otherwise enforce even division
      if num_kv_heads is not None:
          if rank is None:
              rank = dist.get_rank()
          my_slices = (num_kv_heads // mp_size) + (
              1 if rank < (num_kv_heads % mp_size) else 0
          )
          return total_size * my_slices // num_kv_heads
      else:
          if total_size % mp_size == 0:
              return total_size // mp_size
          else:
              assert (
                  False
              ), f"Number of attention heads ({total_size}) must be divisible by mp_size ({mp_size})"
  
imports: from deepspeed import comm
label_desc: Use image node_id 3 for calling a global function with example usage: get_shard_size(total_size, mp_size, rank) and returns: unknown, unknown