class_docstr: Embedding with the token distributions as a weighted mixture of softmaxes.
  
  Expressing the token distributions in this way improves expressiveness and
  enables the matrix of token probabilities given all contexts to be high rank.
  
  The vocab embedding is the same as the default, which is just a simple
  embedding.
  
  See https://arxiv.org/pdf/1711.03953.pdf for more details.
docstr: Configurable embedding for the vocabulary.
  
  Most of the arguments get passed to `mtf.layers.embedding_weights`.
  
  Args:
    mesh: the mesh used to layout the tensors.
    vocab_dim: the dimension corresponding to vocabulary.
    output_dim: the dimension corresponding to the model
      hidden states.
    variable_dtype: the datatype information for the
      variables used in the embedding tensors.
    name: a name to base variable names off of.
    ensemble_dim: the dimension used for ensembling.
      Absolutely no guarantees that this code will work with ensembling.
    num_softmaxes: a positive int, the number of components to use in the
      mixture.
docstr_tok: ["Configurable","embedding","for","the","vocabulary",".","Most","of","the","arguments","get","passed","to","`","mtf.layers.embedding_weights","`",".","Args",":","mesh",":","the","mesh","used","to","layout","the","tensors",".","vocab_dim",":","the","dimension","corresponding","to","vocabulary",".","output_dim",":","the","dimension","corresponding","to","the","model","hidden","states",".","variable_dtype",":","the","datatype","information","for","the","variables","used","in","the","embedding","tensors",".","name",":","a","name","to","base","variable","names","off","of",".","ensemble_dim",":","the","dimension","used","for","ensembling",".","Absolutely","no","guarantees","that","this","code","will","work","with","ensembling",".","num_softmaxes",":","a","positive","int",",","the","number","of","components","to","use","in","the","mixture","."]
code: def __init__(
      self,
      mesh: mtf.Mesh,
      vocab_dim: mtf.Dimension,
      output_dim: mtf.Dimension,
      variable_dtype: mtf.VariableDType,
      name: str,
      ensemble_dim: mtf.Dimension,
      num_softmaxes: int = gin.REQUIRED,
  ):
      """Configurable embedding for the vocabulary.
  
      Most of the arguments get passed to `mtf.layers.embedding_weights`.
  
      Args:
        mesh: the mesh used to layout the tensors.
        vocab_dim: the dimension corresponding to vocabulary.
        output_dim: the dimension corresponding to the model
          hidden states.
        variable_dtype: the datatype information for the
          variables used in the embedding tensors.
        name: a name to base variable names off of.
        ensemble_dim: the dimension used for ensembling.
          Absolutely no guarantees that this code will work with ensembling.
        num_softmaxes: a positive int, the number of components to use in the
          mixture.
      """
      self._vocab_dim = vocab_dim
      self._output_dim = output_dim
      self._copy_output_dim = mtf.Dimension(
          "_{}_copy".format(output_dim.name), output_dim.size
      )
      self._components_dim = mtf.Dimension(
          "softmax_components", num_softmaxes
      )
  
      self._embedding_weights = mtf.layers.embedding_weights(
          mesh=mesh,
          vocab_dim=vocab_dim,
          output_dim=output_dim,
          variable_dtype=variable_dtype,
          name="{}_embedding_weights".format(name),
          ensemble_dim=ensemble_dim,
      )
      self._mixture_weights = mtf.layers.embedding_weights(
          mesh=mesh,
          vocab_dim=self._components_dim,
          output_dim=output_dim,
          variable_dtype=variable_dtype,
          name="{}_mixture_weights".format(name),
          ensemble_dim=ensemble_dim,
      )
      self._context_weights = mtf.layers.embedding_weights(
          mesh=mesh,
          vocab_dim=self._copy_output_dim,
          output_dim=output_dim,
          variable_dtype=variable_dtype,
          name="{}_context_weights".format(name),
          ensemble_dim=([ensemble_dim] if ensemble_dim else [])
          + [self._components_dim],
      )
  
imports: import gin
  import mesh_tensorflow
  from mesh_tensorflow.transformer import transformer
  import tensorflow.compat.v1
label_desc: Use image node_id 1 to create a new MixtureOfSoftmaxes object from inherited base classes: object with example: obj = MixtureOfSoftmaxes(mesh, vocab_dim, output_dim, variable_dtype, name, ensemble_dim, num_softmaxes)