class_docstr: A set of parameters used for (multihead) attention.
docstr: Initialize attention projection matrices.
docstr_tok: ["Initialize","attention","projection","matrices","."]
code: def init_weights(self):
      """Initialize attention projection matrices."""
      if mtf.layers.unit_scaling_convention():
          init = tf.random_normal_initializer(stddev=1.0)
          q_init = init
          kv_init = init
          o_init = init
      else:
          stddev = self.query_input_dim.size**-0.5
          if self.fold_scaling_into_initializer:
              stddev *= self.key_dim.size**-0.5
          q_init = tf.random_normal_initializer(stddev=stddev)
          kv_init = tf.random_normal_initializer(
              stddev=self.memory_input_dim.size**-0.5
          )
          o_init = tf.random_normal_initializer(
              stddev=mtf.Shape(
                  self.query_heads_dims + [self.value_dim]
              ).size
              ** -0.5
          )
  
      # Toggle producing wq, wv, wk which are not needed for the ExpertsAttention
      if self.make_attention_vars:
          if not self.no_query:
              self.wq = mtf.get_variable(
                  self.mesh,
                  "q",
                  self.q_shape,
                  initializer=q_init,
                  dtype=self.variable_dtype,
              )
          if self.shared_kv:
              self.wkv = mtf.get_variable(
                  self.mesh,
                  "kv",
                  self.k_shape,
                  initializer=kv_init,
                  dtype=self.variable_dtype,
              )
          else:
              self.wk = mtf.get_variable(
                  self.mesh,
                  "k",
                  self.k_shape,
                  initializer=kv_init,
                  dtype=self.variable_dtype,
              )
              self.wv = mtf.get_variable(
                  self.mesh,
                  "v",
                  self.v_shape,
                  initializer=kv_init,
                  dtype=self.variable_dtype,
              )
      self.wo = mtf.get_variable(
          self.mesh,
          "o",
          self.o_shape,
          initializer=o_init,
          dtype=self.variable_dtype,
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import mesh_tensorflow
  import tensorflow.compat.v1
label_desc: Use image node_id 2 for calling the AttentionParams obj's underlying member method code with example usage: obj.init_weights() without return types