class_docstr: Embedding that uses Mixtape in computing logits.
  
  Expressing the token distributions in this way improves expressiveness and
  enables the matrix of token probabilities given all contexts to be high rank.
  
  Mixtape has the advantage of added efficiency over other methods such as
  mixture of softmax.
  
  The vocab embedding is the same as the default, which just a simple embedding.
  
  See
  https://papers.nips.cc/paper/9723-mixtape-breaking-the-softmax-bottleneck-efficiently.pdf
  for more details.
docstr: Configurable embedding for the vocabulary.
  
  Most of the arguments get passed to `mtf.layers.embedding_weights`.
  
  Mixtape shares gates for low frequency tokens to improve efficiency. Since
  our vocabs are sorted in decreasing order of frequency with sentinels
  appended to the end, we need to do a little trick to ensure that the
  sentinels are treated as high frequency. If you want to treat the sentinels
  as low frequency tokens, then pass in zero for `extra_ids`.
  
  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.
    extra_ids: a non-negative integer, the number of sentinels at the end of
      the vocab.
    dropout_rate: a float between 0 and 1, the rate to use for dropout.
    gate_embedding_size: a positive integer, the size to use for embedding for
      the gates. It is usually chosen to be much smaller than d_model.
    frequent_token_fraction: a float between 0 and 1, what fraction of tokens
      to consider as high frequency and not share gates for.
    noise_std_dev: a non-negative float, the standard deviation of the
      Gaussian noise to add to the pre-activation priors.
docstr_tok: ["Configurable","embedding","for","the","vocabulary",".","Most","of","the","arguments","get","passed","to","`","mtf.layers.embedding_weights","`",".","Mixtape","shares","gates","for","low","frequency","tokens","to","improve","efficiency",".","Since","our","vocabs","are","sorted","in","decreasing","order","of","frequency","with","sentinels","appended","to","the","end",",","we","need","to","do","a","little","trick","to","ensure","that","the","sentinels","are","treated","as","high","frequency",".","If","you","want","to","treat","the","sentinels","as","low","frequency","tokens",",","then","pass","in","zero","for","`","extra_ids","`",".","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",".","extra_ids",":","a","non-negative","integer",",","the","number","of","sentinels","at","the","end","of","the","vocab",".","dropout_rate",":","a","float","between","0","and","1",",","the","rate","to","use","for","dropout",".","gate_embedding_size",":","a","positive","integer",",","the","size","to","use","for","embedding","for","the","gates",".","It","is","usually","chosen","to","be","much","smaller","than","d_model",".","frequent_token_fraction",":","a","float","between","0","and","1",",","what","fraction","of","tokens","to","consider","as","high","frequency","and","not","share","gates","for",".","noise_std_dev",":","a","non-negative","float",",","the","standard","deviation","of","the","Gaussian","noise","to","add","to","the","pre-activation","priors","."]
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,
      extra_ids: int = 0,
      dropout_rate: float = 0.0,
      gate_embedding_size: int = gin.REQUIRED,
      frequent_token_fraction: float = 0.1,
      noise_std_dev: float = 0.0,
  ):
      """Configurable embedding for the vocabulary.
  
      Most of the arguments get passed to `mtf.layers.embedding_weights`.
  
      Mixtape shares gates for low frequency tokens to improve efficiency. Since
      our vocabs are sorted in decreasing order of frequency with sentinels
      appended to the end, we need to do a little trick to ensure that the
      sentinels are treated as high frequency. If you want to treat the sentinels
      as low frequency tokens, then pass in zero for `extra_ids`.
  
      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.
        extra_ids: a non-negative integer, the number of sentinels at the end of
          the vocab.
        dropout_rate: a float between 0 and 1, the rate to use for dropout.
        gate_embedding_size: a positive integer, the size to use for embedding for
          the gates. It is usually chosen to be much smaller than d_model.
        frequent_token_fraction: a float between 0 and 1, what fraction of tokens
          to consider as high frequency and not share gates for.
        noise_std_dev: a non-negative float, the standard deviation of the
          Gaussian noise to add to the pre-activation priors.
      """
      self._extra_ids = extra_ids
      self._dropout_rate = dropout_rate
      self._noise_std_dev = noise_std_dev
      self._mesh = mesh
      self._vocab_dim = vocab_dim
      self._frequent_vocab_dim = mtf.Dimension(
          vocab_dim.name, int(frequent_token_fraction * vocab_dim.size)
      )
      self._rare_vocab_dim = mtf.Dimension(
          vocab_dim.name, vocab_dim.size - self._frequent_vocab_dim.size
      )
      self._output_dim = output_dim
      self._copy_output_dim = mtf.Dimension(
          "_{}_copy".format(output_dim.name), output_dim.size
      )
      self._pre_gates_dim = mtf.Dimension("gates", 3)
      self._gates_dim = mtf.Dimension("gates", 4)
      self._gate_embedding_dim = mtf.Dimension(
          "gate_embedding", gate_embedding_size
      )
  
      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,
      )
      ensemble_dims = [ensemble_dim] if ensemble_dim else []
      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_dims + [self._gates_dim],
      )
      self._context_weights_bias = mtf.get_variable(
          mesh,
          name="{}_context_weights_bias".format(name),
          shape=mtf.Shape(
              ensemble_dims + [self._gates_dim, output_dim]
          ),
          dtype=variable_dtype,
          initializer=tf.zeros_initializer(),
      )
  
      self._prior_weights = mtf.layers.embedding_weights(
          mesh=mesh,
          vocab_dim=self._gate_embedding_dim,
          output_dim=output_dim,
          variable_dtype=variable_dtype,
          name="{}_prior_weights".format(name),
          ensemble_dim=ensemble_dims + [self._pre_gates_dim],
      )
      self._prior_weights_bias = mtf.get_variable(
          mesh,
          name="{}_prior_weights_bias".format(name),
          shape=mtf.Shape(
              ensemble_dims
              + [self._pre_gates_dim, self._gate_embedding_dim]
          ),
          dtype=variable_dtype,
          initializer=tf.zeros_initializer(),
      )
      self._prior_vocab_vector = mtf.get_variable(
          mesh,
          name="{}_prior_vocab_vector".format(name),
          shape=mtf.Shape(
              ensemble_dims
              + [self._frequent_vocab_dim, self._gate_embedding_dim]
          ),
          dtype=variable_dtype,
          initializer=tf.random_normal_initializer(),
      )
      self._prior_gates_vector = mtf.get_variable(
          mesh,
          name="{}_prior_gates_vector".format(name),
          shape=mtf.Shape(
              ensemble_dims + [self._pre_gates_dim, output_dim]
          ),
          dtype=variable_dtype,
          initializer=tf.random_normal_initializer(),
      )
      self._prior_bias = mtf.get_variable(
          mesh,
          name="{}_prior_bias".format(name),
          shape=mtf.Shape(
              ensemble_dims
              + [self._frequent_vocab_dim, self._pre_gates_dim]
          ),
          dtype=variable_dtype,
          initializer=tf.random_normal_initializer(),
      )
  
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 Mixtape object from inherited base classes: object with example: obj = Mixtape(mesh, vocab_dim, output_dim, variable_dtype, name, ensemble_dim, extra_ids, dropout_rate, gate_embedding_size, frequent_token_fraction, noise_std_dev)