class_docstr: Multi-head self-attention layer with the Fixup initialization.
class_docstr_tok: ["Multi-head","self-attention","layer","with","the","Fixup","initialization","."]
code: def make_params(self, context):
      if self.num_heads == 1:
          query_heads_dims = None
          memory_heads_dims = None
      elif self.num_memory_heads == 0:
          query_heads_dims = [mtf.Dimension("heads", self.num_heads)]
          memory_heads_dims = query_heads_dims
      elif self.num_memory_heads == 1:
          query_heads_dims = [mtf.Dimension("heads", self.num_heads)]
          memory_heads_dims = None
      else:
          if self.num_heads % self.num_memory_heads != 0:
              raise ValueError("num_memory_heads must divide num_heads")
          memory_heads_dims = [
              mtf.Dimension("heads", self.num_memory_heads)
          ]
          query_heads_dims = memory_heads_dims + [
              mtf.Dimension(
                  "query_heads", self.num_heads // self.num_memory_heads
              )
          ]
  
      return AttentionParamsFixup(
          context.mesh,
          query_input_dim=context.model.model_dim,
          memory_input_dim=context.model.model_dim,
          output_dim=context.model.model_dim,
          key_dim=self.kv_dim,
          value_dim=self.kv_dim,
          query_heads_dims=query_heads_dims,
          memory_heads_dims=memory_heads_dims,
          variable_dtype=context.variable_dtype,
          shared_kv=self.shared_kv,
          ensemble_dim=context.model.ensemble_dim,
          combine_dims=self.combine_dims,
          keep_query_heads_dims=self.keep_query_heads_dims,
          fold_scaling_into_initializer=self.fold_scaling_into_initializer,
          num_blocks=self.num_blocks,
          default_init=self.default_init,
          init_distribution=self.init_distribution,
      )
  
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 SelfAttentionFixup obj's underlying member method code with example usage: obj.make_params(context) and returns: AttentionParamsFixup