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: Function called by mtf transformer to get the logits.
  
  Note that we are taking the log of a mixture of softmaxes. The logits will
  then go through a softmax. This could potentially run into numerical
  stability issues. If that happens, try setting the activation_dtype to
  float32.
  
  Args:
    hidden: hidden model states of the final decoder layer.
    context: the context used for the call to the
      transformer.
  
  Returns:
    The logits.
docstr_tok: ["Function","called","by","mtf","transformer","to","get","the","logits",".","Note","that","we","are","taking","the","log","of","a","mixture","of","softmaxes",".","The","logits","will","then","go","through","a","softmax",".","This","could","potentially","run","into","numerical","stability","issues",".","If","that","happens",",","try","setting","the","activation_dtype","to","float32",".","Args",":","hidden",":","hidden","model","states","of","the","final","decoder","layer",".","context",":","the","context","used","for","the","call","to","the","transformer",".","Returns",":","The","logits","."]
code: def hidden_to_logits(
      self, hidden: mtf.Tensor, context: transformer.Context
  ) -> mtf.Tensor:
      """Function called by mtf transformer to get the logits.
  
      Note that we are taking the log of a mixture of softmaxes. The logits will
      then go through a softmax. This could potentially run into numerical
      stability issues. If that happens, try setting the activation_dtype to
      float32.
  
      Args:
        hidden: hidden model states of the final decoder layer.
        context: the context used for the call to the
          transformer.
  
      Returns:
        The logits.
      """
      del context
      hidden *= self._output_dim.size**-0.5
  
      component_prior_logits = mtf.einsum(
          [hidden, self._mixture_weights],
          reduced_dims=[self._output_dim],
      )
  
      component_contexts = mtf.einsum(
          [
              mtf.rename_dimension(
                  hidden,
                  self._output_dim.name,
                  self._copy_output_dim.name,
              ),
              self._context_weights,
          ],
          reduced_dims=[self._copy_output_dim],
      )
      component_contexts = mtf.tanh(component_contexts)
      component_logits = mtf.einsum(
          [component_contexts, self._embedding_weights],
          reduced_dims=[self._output_dim],
      )
  
      component_prior_logits = mtf.log_softmax(
          component_prior_logits, reduced_dim=self._components_dim
      )
      component_logits = mtf.log_softmax(
          component_logits, reduced_dim=self._vocab_dim
      )
  
      logits = component_prior_logits + component_logits
      logits = mtf.reduce_logsumexp(
          logits, reduced_dim=self._components_dim
      )
      return logits
  
imports: import gin
  import mesh_tensorflow
  from mesh_tensorflow.transformer import transformer
  import tensorflow.compat.v1
label_desc: Use image node_id 3 for calling the MixtureOfSoftmaxes obj's underlying member method code with example usage: obj.hidden_to_logits(hidden, context) and returns: logits