class_docstr: Vocab embedding implementing the adaptive softmax.
  
  The adaptive softmax was first introduced in this paper
  (https://arxiv.org/abs/1609.04309). Note that this implementation is actually
  most similar to the adaptive vocab embeddings in
  https://arxiv.org/abs/1809.10853 as it supports having different embedding
  sizes for different clusters.
  
  The adaptive softmax works by factorizing the traditional softmax over
  multiple clusters:
    p(v|h) = p(v|c,h) p(c|h),
  where both probability distributions take the form of a softmax.
  
  Further speed up is achieved by putting the class containing the most
  frequently occurring tokens in the "head" cluster. Essentially, those tokens
  are included as "classes" in the p(c|h) softmax. Thus computing their
  probabilities requires only single softmax evaluation.
  
  This implementation differs from vocab_embeddings.AdaptiveVocabEmbedding. That
  implementation only supports variable embeddings sizes across clusters. This
  implementation also supports the adaptive softmax.
  
  A few conditions must be met in order to use this vocab:
    - Unitransformer.shared_embedding_and_softmax_weights = True.
    - If training, then
      Unitranformer.loss_fn = adaptive_softmax.adaptive_softmax_loss_fn.
    - Label smoothing is not supported and will be ignored silently.
    - loss_on_targets_only is not supported and will be ignored silently.
class_docstr_tok: ["Vocab","embedding","implementing","the","adaptive","softmax",".","The","adaptive","softmax","was","first","introduced","in","this","paper","(","https",":","\/\/arxiv.org\/abs\/1609.04309",")",".","Note","that","this","implementation","is","actually","most","similar","to","the","adaptive","vocab","embeddings","in","https",":","\/\/arxiv.org\/abs\/1809.10853","as","it","supports","having","different","embedding","sizes","for","different","clusters",".","The","adaptive","softmax","works","by","factorizing","the","traditional","softmax","over","multiple","clusters",":","p","(","v|h",")","=","p","(","v|c",",","h",")","p","(","c|h",")",",","where","both","probability","distributions","take","the","form","of","a","softmax",".","Further","speed","up","is","achieved","by","putting","the","class","containing","the","most","frequently","occurring","tokens","in","the","``","head","''","cluster",".","Essentially",",","those","tokens","are","included","as","``","classes","''","in","the","p","(","c|h",")","softmax",".","Thus","computing","their","probabilities","requires","only","single","softmax","evaluation",".","This","implementation","differs","from","vocab_embeddings.AdaptiveVocabEmbedding",".","That","implementation","only","supports","variable","embeddings","sizes","across","clusters",".","This","implementation","also","supports","the","adaptive","softmax",".","A","few","conditions","must","be","met","in","order","to","use","this","vocab",":","-","Unitransformer.shared_embedding_and_softmax_weights","=","True",".","-","If","training",",","then","Unitranformer.loss_fn","=","adaptive_softmax.adaptive_softmax_loss_fn",".","-","Label","smoothing","is","not","supported","and","will","be","ignored","silently",".","-","loss_on_targets_only","is","not","supported","and","will","be","ignored","silently","."]
docstr: Function called by mtf transformer to get the logits.
  
  The benefit from the adaptive softmax comes from not having to compute the
  logits over all of the vocab during training. Thus, we use the somewhat
  hacky solution of returning the hidden states during training and then using
  them to compute the loss in a custom loss function.
  
  When not training, this method will be true to its name as return the
  logits corresponding to the hidden state.
  
  Args:
    hidden: an mtf.Tensor, hidden model states of the final decoder layer.
    context: a transformer.Context, the context used for the call to the
      transformer.
  
  Returns:
    an mtf.Tensor
docstr_tok: ["Function","called","by","mtf","transformer","to","get","the","logits",".","The","benefit","from","the","adaptive","softmax","comes","from","not","having","to","compute","the","logits","over","all","of","the","vocab","during","training",".","Thus",",","we","use","the","somewhat","hacky","solution","of","returning","the","hidden","states","during","training","and","then","using","them","to","compute","the","loss","in","a","custom","loss","function",".","When","not","training",",","this","method","will","be","true","to","its","name","as","return","the","logits","corresponding","to","the","hidden","state",".","Args",":","hidden",":","an","mtf.Tensor",",","hidden","model","states","of","the","final","decoder","layer",".","context",":","a","transformer.Context",",","the","context","used","for","the","call","to","the","transformer",".","Returns",":","an","mtf.Tensor"]
code: def hidden_to_logits(
      self, hidden: mtf.Tensor, context: transformer.Context
  ) -> mtf.Tensor:
      """Function called by mtf transformer to get the logits.
  
      The benefit from the adaptive softmax comes from not having to compute the
      logits over all of the vocab during training. Thus, we use the somewhat
      hacky solution of returning the hidden states during training and then using
      them to compute the loss in a custom loss function.
  
      When not training, this method will be true to its name as return the
      logits corresponding to the hidden state.
  
      Args:
        hidden: an mtf.Tensor, hidden model states of the final decoder layer.
        context: a transformer.Context, the context used for the call to the
          transformer.
  
      Returns:
        an mtf.Tensor
      """
      if context.mode == tf_estimator.ModeKeys.TRAIN:
          return hidden
      else:
          return self._hidden_to_logits(hidden, context)
  
imports: import math
  from typing import Dict, Sequence, Union
  import gin
  import mesh_tensorflow
  from mesh_tensorflow.transformer import transformer
  from mesh_tensorflow.transformer import vocab_embeddings
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
label_desc: Use image node_id 4 for calling the AdaptiveSoftmaxVocabEmbedding obj's underlying member method code with example usage: obj.hidden_to_logits(hidden, context) and returns: hidden, self