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: Returns the loss without computing a softmax over the entire vocab.
docstr_tok: ["Returns","the","loss","without","computing","a","softmax","over","the","entire","vocab","."]
code: def compute_loss(
      self,
      decoder: transformer.Unitransformer,
      hidden: mtf.Tensor,
      targets: mtf.Tensor,
      context: transformer.Context,
  ) -> mtf.Tensor:
      """Returns the loss without computing a softmax over the entire vocab."""
      loss = 0
      tail_cluster_masks = []
      for cluster in self._tail_clusters:
          cluster_mask = cluster.get_cluster_mask(targets)
          tail_cluster_masks.append(cluster_mask)
  
          if cluster.length_projection_factor == 1:
              targets_in_cluster = mtf.where(cluster_mask, targets, 0)
              hidden_in_cluster = mtf.where(cluster_mask, hidden, 0)
          else:
              # TODO(mmatena): Unfold the batch dim to get a super long sequence dim
              # to reduce the risk of overflowing the projection.
              proj_to_cluster_len = (
                  cluster.get_project_to_cluster_length(
                      cluster_mask, dtype=targets.dtype
                  )
              )
              targets_in_cluster = mtf.einsum(
                  [proj_to_cluster_len, targets],
                  reduced_dims=[
                      targets.shape.get_dim_by_name("length")
                  ],
              )
              hidden_in_cluster = mtf.einsum(
                  [mtf.cast(proj_to_cluster_len, hidden.dtype), hidden],
                  reduced_dims=[hidden.shape.get_dim_by_name("length")],
              )
  
          loss += cluster.compute_loss(
              decoder, hidden_in_cluster, targets_in_cluster, context
          )
  
      tail_clusters_dim = mtf.Dimension(
          "tail_clusters", len(tail_cluster_masks)
      )
      tail_node_targets = mtf.reduce_sum(
          mtf.stack(
              [
                  (self._head_cluster.end_token_id + i)
                  * mtf.cast(mask, targets.dtype)
                  for i, mask in enumerate(tail_cluster_masks)
              ],
              tail_clusters_dim.name,
          ),
          reduced_dim=tail_clusters_dim,
      )
      head_targets = mtf.where(
          mtf.cast(tail_node_targets, tf.bool),
          tail_node_targets,
          targets,
      )
      loss += self._head_cluster.compute_loss(
          decoder, hidden, head_targets, context
      )
  
      return loss
  
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 6 for calling the AdaptiveSoftmaxVocabEmbedding obj's underlying member method code with example usage: obj.compute_loss(decoder, hidden, targets, context) and returns: loss