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_weights(self):
      o_init_fixup = tf.initializers.zeros()
  
      # Since tf.initializers.variance_scaling returns sqrt(3 * scale / n), (note
      # that scale is inside sqrt), we need to square the scale factor. Hence the
      # exponent is -1/3 instead of -1/6 as described in the class docstring.
      if self.default_init == "glorot":
          init_fixup = tf.initializers.variance_scaling(
              mode="fan_avg",
              distribution=self.init_distribution,
              scale=math.pow(self.num_blocks, -1.0 / 3),
          )
      elif self.default_init == "he":
          init_fixup = tf.initializers.variance_scaling(
              mode="fan_in",
              distribution=self.init_distribution,
              scale=2 * math.pow(self.num_blocks, -1.0 / 3),
          )
      else:
          raise ValueError(
              (
                  "Unsupported default initialization. Only 'glorot' and 'he'"
                  " initializations are supported."
              )
          )
  
      if not self.no_query:
          self.wq = mtf.get_variable(
              self.mesh,
              "q_fixup",
              self.q_shape,
              initializer=init_fixup,
              dtype=self.variable_dtype,
          )
  
      if self.shared_kv:
          self.wkv = mtf.get_variable(
              self.mesh,
              "kv_fixup",
              self.k_shape,
              initializer=init_fixup,
              dtype=self.variable_dtype,
          )
      else:
          self.wk = mtf.get_variable(
              self.mesh,
              "k_fixup",
              self.k_shape,
              initializer=init_fixup,
              dtype=self.variable_dtype,
          )
  
          self.wv = mtf.get_variable(
              self.mesh,
              "v_fixup",
              self.v_shape,
              initializer=init_fixup,
              dtype=self.variable_dtype,
          )
  
      self.wo = mtf.get_variable(
          self.mesh,
          "o_fixup",
          self.o_shape,
          initializer=o_init_fixup,
          dtype=self.variable_dtype,
      )
  
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 2 for calling the AttentionParamsFixup obj's underlying member method code with example usage: obj.init_weights() without return types