code: 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)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from tensor2tensor.layers import common_attention
  from tensor2tensor.layers import common_layers
  from tensor2tensor.models import transformer
  from tensor2tensor.utils import registry
  from tensor2tensor.utils import t2t_model
  import tensorflow.compat.v1
  from tensorflow.python.ops import inplace_ops
label_desc: Use image node_id 1 for calling a global function with example usage: _capped_double_heads(num_heads, cap) and returns: max