class_docstr: A set of parameters used for (multihead) attention.
docstr: Create attention parameters.
  
  combine_dims is a hack for faster execution.  The heads and key/value
  dimensions are combined in the variables and the computation.  The hack
  would not be necessary if XLA optimized einsum properly.
  
  Args:
    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
    no_query: a boolean
    combine_dims: a boolean
    ensemble_dim: an optional Dimension
    keep_query_heads_dims: a boolean, if true keep the query_heads_dims in the
      output.
    fold_scaling_into_initializer: a boolean
    make_attention_vars: a boolean, whether to make the attention variables.
      This is typically True. Only set to False for ExpertsAttention which
      creates variables inside the moe.MoE1D-call.
docstr_tok: ["Create","attention","parameters",".","combine_dims","is","a","hack","for","faster","execution",".","The","heads","and","key\/value","dimensions","are","combined","in","the","variables","and","the","computation",".","The","hack","would","not","be","necessary","if","XLA","optimized","einsum","properly",".","Args",":","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","no_query",":","a","boolean","combine_dims",":","a","boolean","ensemble_dim",":","an","optional","Dimension","keep_query_heads_dims",":","a","boolean",",","if","true","keep","the","query_heads_dims","in","the","output",".","fold_scaling_into_initializer",":","a","boolean","make_attention_vars",":","a","boolean",",","whether","to","make","the","attention","variables",".","This","is","typically","True",".","Only","set","to","False","for","ExpertsAttention","which","creates","variables","inside","the","moe.MoE1D-call","."]
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,
      no_query=False,
      combine_dims=True,
      ensemble_dim=None,
      keep_query_heads_dims=False,
      fold_scaling_into_initializer=True,
      make_attention_vars=True,
  ):
      """Create attention parameters.
  
      combine_dims is a hack for faster execution.  The heads and key/value
      dimensions are combined in the variables and the computation.  The hack
      would not be necessary if XLA optimized einsum properly.
  
      Args:
        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
        no_query: a boolean
        combine_dims: a boolean
        ensemble_dim: an optional Dimension
        keep_query_heads_dims: a boolean, if true keep the query_heads_dims in the
          output.
        fold_scaling_into_initializer: a boolean
        make_attention_vars: a boolean, whether to make the attention variables.
          This is typically True. Only set to False for ExpertsAttention which
          creates variables inside the moe.MoE1D-call.
      """
      if shared_kv and key_dim != value_dim:
          raise ValueError("shared_kv requires key_dim == value_dim")
      self.mesh = mesh
      self.query_input_dim = query_input_dim
      self.memory_input_dim = memory_input_dim
      self.output_dim = output_dim
      self.key_dim = key_dim
      self.value_dim = value_dim
      self.query_heads_dims = query_heads_dims or []
      self.memory_heads_dims = memory_heads_dims or []
      self.variable_dtype = variable_dtype
      self.shared_kv = shared_kv
      self.no_query = no_query
      self.combine_dims = combine_dims
      self.keep_query_heads_dims = keep_query_heads_dims
      self.fold_scaling_into_initializer = fold_scaling_into_initializer
      self.make_attention_vars = make_attention_vars
  
      if combine_dims:
          self.q_shape = [query_input_dim, _combined_dim(self.q_dims)]
          self.k_shape = [memory_input_dim, _combined_dim(self.k_dims)]
          self.v_shape = [memory_input_dim, _combined_dim(self.v_dims)]
          self.o_shape = [_combined_dim(self.o_dims), output_dim]
      else:
          self.q_shape = [query_input_dim] + self.q_dims
          self.k_shape = [memory_input_dim] + self.k_dims
          self.v_shape = [memory_input_dim] + self.v_dims
          self.o_shape = self.o_dims + [output_dim]
      if ensemble_dim:
          self.q_shape = [ensemble_dim] + self.q_shape
          self.k_shape = [ensemble_dim] + self.k_shape
          self.v_shape = [ensemble_dim] + self.v_shape
          self.o_shape = [ensemble_dim] + self.o_shape
  
      self.init_weights()
  
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 1 to create a new AttentionParams object from inherited base classes: object with example: obj = AttentionParams(mesh, query_input_dim, memory_input_dim, output_dim, key_dim, value_dim, query_heads_dims, memory_heads_dims, variable_dtype, shared_kv, no_query, combine_dims, ensemble_dim, keep_query_heads_dims, fold_scaling_into_initializer, make_attention_vars)