def _capped_double_heads(num_heads, cap=16):
  """Calculate the number of heads for the attention layers with more heads.

  The number of heads will be twice the normal amount (num_heads), until it
  reaches |cap| heads.

  Args:
    num_heads: the num_heads hparam for the model.
    cap: the maximum number of heads |num_heads| will be doubled to.

  Returns:
    The number of heads for the attention layers that have more heads.
  """
  return max(min(num_heads * 2, cap), num_heads)