class_docstr: Create attention parameters with Fixup initialization.
  
  See class docstring for DenseReluDenseFixup for details.
  
  For SelfAttention layer, m = 4, i.e., 4 weight matrix multiplications. See
  https://github.com/hongyi-zhang/Fixup/issues/8#issuecomment-505750941.
  So the scaling factor for SelfAttention layer is num_blocks**(-1/6).
  
  Attributes:
    mesh: a Mesh
    query_input_dim: a Dimension
    memory_input_dim: a Dimension
    output_dim: a Dimension
    key_dim: a Dimension
    value_dim: a Dimension
    query_heads_dims: a list of Dimension
    memory_heads_dims: a list of Dimension
    variable_dtype: a mtf.VariableDType
    shared_kv: a boolean
    fold_scaling_into_initializer: a boolean
    num_blocks: an integer specifying the number of TransformerLayer objects.
      For a vanilla Transformer model with 12 encoder layers and 12 decoder
      layers, the number of blocks is 2 * 12 + 3 * 12 = 60 where each encoder
      layer has 2 blocks, SelfAttention and Feedforward block and decoder
      additionally has the encoder-decoder attention block.
    o_init_fixup: a tf.initializer for the self.wo.
    init_fixup: a tf.initializer for the self.wq, self.wk, self.wv and self.wkv.
class_docstr_tok: ["Create","attention","parameters","with","Fixup","initialization",".","See","class","docstring","for","DenseReluDenseFixup","for","details",".","For","SelfAttention","layer",",","m","=","4",",","i.e.",",","4","weight","matrix","multiplications",".","See","https",":","\/\/github.com\/hongyi-zhang\/Fixup\/issues\/8","#","issuecomment-505750941",".","So","the","scaling","factor","for","SelfAttention","layer","is","num_blocks","*","*","(","-1\/6",")",".","Attributes",":","mesh",":","a","Mesh","query_input_dim",":","a","Dimension","memory_input_dim",":","a","Dimension","output_dim",":","a","Dimension","key_dim",":","a","Dimension","value_dim",":","a","Dimension","query_heads_dims",":","a","list","of","Dimension","memory_heads_dims",":","a","list","of","Dimension","variable_dtype",":","a","mtf.VariableDType","shared_kv",":","a","boolean","fold_scaling_into_initializer",":","a","boolean","num_blocks",":","an","integer","specifying","the","number","of","TransformerLayer","objects",".","For","a","vanilla","Transformer","model","with","12","encoder","layers","and","12","decoder","layers",",","the","number","of","blocks","is","2","*","12","+","3","*","12","=","60","where","each","encoder","layer","has","2","blocks",",","SelfAttention","and","Feedforward","block","and","decoder","additionally","has","the","encoder-decoder","attention","block",".","o_init_fixup",":","a","tf.initializer","for","the","self.wo",".","init_fixup",":","a","tf.initializer","for","the","self.wq",",","self.wk",",","self.wv","and","self.wkv","."]
code: def __init__(
      self,
      mesh,
      query_input_dim,
      memory_input_dim,
      output_dim,
      key_dim,
      value_dim,
      query_heads_dims,
      memory_heads_dims,
      variable_dtype,
      shared_kv=False,
      fold_scaling_into_initializer=False,
      num_blocks=None,
      default_init="he",
      init_distribution="uniform",
      **kwargs,
  ):
      self.num_blocks = num_blocks
      self.default_init = default_init
      self.init_distribution = init_distribution
  
      if mtf.layers.unit_scaling_convention():
          raise ValueError(
              "Fixup initialization is not compatible with unit scaling convention."
          )
  
      if fold_scaling_into_initializer:
          raise ValueError(
              "Fixup initialization is not compatible with "
              "`fold_scaling_into_initializer."
          )
  
      super(AttentionParamsFixup, self).__init__(
          mesh,
          query_input_dim,
          memory_input_dim,
          output_dim,
          key_dim,
          value_dim,
          query_heads_dims,
          memory_heads_dims,
          variable_dtype,
          shared_kv=shared_kv,
          fold_scaling_into_initializer=fold_scaling_into_initializer,
          **kwargs,
      )
  
imports: import math
  import gin
  import mesh_tensorflow
  from mesh_tensorflow.transformer import attention
  from mesh_tensorflow.transformer import transformer
  from mesh_tensorflow.transformer import transformer_layers
  import tensorflow.compat.v1
label_desc: Use image node_id 1 to create a new AttentionParamsFixup object from inherited base classes: attention with example: obj = AttentionParamsFixup(mesh, query_input_dim, memory_input_dim, output_dim, key_dim, value_dim, query_heads_dims, memory_heads_dims, variable_dtype, shared_kv, fold_scaling_into_initializer, num_blocks, default_init, init_distribution)