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: Configurable embedding for the vocabulary.
  
  Most of the arguments get passed to `mtf.layers.embedding_weights`.
  
  The clustering parameters are specified by the `clusters` argument. It is a
  list of dicts with keys:
    - token_count: The number of tokens in the cluster.
    - embedding_size: (optional) The hidden dimension size of the cluster's
      embedding. Defaults to the model dimension size.
    - length_projection_factor: (optional) Since MTF can't handle variable
      length dimensions, we project from the sequence length dimension to a
      dimension of size length_projection_factor * sequence_length during
      training. This can save compute time and resources if the cluster has
      many tokens that appear infrequently. If all of the tokens belonging to
      the cluster cannot fit within this reduced dimension, some will be
      discarded and ignored for the purposes of computing loss. Defaults 1.
      Ignored for the head (first) cluster and not during training.
  
  The first cluster will become the head cluster.
  
  For example, let's say we have a vocab size of 500k and pass as clusters:
    [
      {"token_count": 50000,  "embedding_size": 1024},
      {"token_count": 100000, "embedding_size": 256},
      {"token_count": 350000, "embedding_size": 64},
    ]
  Then tokens with ids 0 (inclusive) to 50k (exclusive) will be in the first
  cluster with embedding size of 1024, tokens with ids 50k to 150k will be in
  the second cluster with embedding size of 256, and tokens with ids 150k to
  500k will be in the third cluster with embedding size of 64.
  
  Args:
    mesh: a mtf.Mesh, the mesh used to layout the tensors.
    vocab_dim: a mtf.Dimension, the dimension corresponding to vocabulary.
    output_dim: a mtf.Dimension, the dimension corresponding to the model
      hidden states.
    variable_dtype: a mtf.VariableDType, the datatype information for the
      variables used in the embedding tensors.
    name: a string, a name to base variable names off of.
    ensemble_dim: a mtf.Dimension, the dimension used for ensembling.
      Absolutely no guarantees that this code will work with ensembling.
    clusters: a list(dict), specification of the clusters. See above for more
      information.
  
  Raises:
    ValueError: The sum of the token counts across the clusters does not equal
      the vocabulary size or a length_projection_factor is not in the range
      (0, 1].
docstr_tok: ["Configurable","embedding","for","the","vocabulary",".","Most","of","the","arguments","get","passed","to","`","mtf.layers.embedding_weights","`",".","The","clustering","parameters","are","specified","by","the","`","clusters","`","argument",".","It","is","a","list","of","dicts","with","keys",":","-","token_count",":","The","number","of","tokens","in","the","cluster",".","-","embedding_size",":","(","optional",")","The","hidden","dimension","size","of","the","cluster's","embedding",".","Defaults","to","the","model","dimension","size",".","-","length_projection_factor",":","(","optional",")","Since","MTF","ca","n't","handle","variable","length","dimensions",",","we","project","from","the","sequence","length","dimension","to","a","dimension","of","size","length_projection_factor","*","sequence_length","during","training",".","This","can","save","compute","time","and","resources","if","the","cluster","has","many","tokens","that","appear","infrequently",".","If","all","of","the","tokens","belonging","to","the","cluster","can","not","fit","within","this","reduced","dimension",",","some","will","be","discarded","and","ignored","for","the","purposes","of","computing","loss",".","Defaults","1",".","Ignored","for","the","head","(","first",")","cluster","and","not","during","training",".","The","first","cluster","will","become","the","head","cluster",".","For","example",",","let","'s","say","we","have","a","vocab","size","of","500k","and","pass","as","clusters",":","[","{","``","token_count","''",":","50000",",","``","embedding_size","''",":","1024","}",",","{","``","token_count","''",":","100000",",","``","embedding_size","''",":","256","}",",","{","``","token_count","''",":","350000",",","``","embedding_size","''",":","64","}",",","]","Then","tokens","with","ids","0","(","inclusive",")","to","50k","(","exclusive",")","will","be","in","the","first","cluster","with","embedding","size","of","1024",",","tokens","with","ids","50k","to","150k","will","be","in","the","second","cluster","with","embedding","size","of","256",",","and","tokens","with","ids","150k","to","500k","will","be","in","the","third","cluster","with","embedding","size","of","64",".","Args",":","mesh",":","a","mtf.Mesh",",","the","mesh","used","to","layout","the","tensors",".","vocab_dim",":","a","mtf.Dimension",",","the","dimension","corresponding","to","vocabulary",".","output_dim",":","a","mtf.Dimension",",","the","dimension","corresponding","to","the","model","hidden","states",".","variable_dtype",":","a","mtf.VariableDType",",","the","datatype","information","for","the","variables","used","in","the","embedding","tensors",".","name",":","a","string",",","a","name","to","base","variable","names","off","of",".","ensemble_dim",":","a","mtf.Dimension",",","the","dimension","used","for","ensembling",".","Absolutely","no","guarantees","that","this","code","will","work","with","ensembling",".","clusters",":","a","list","(","dict",")",",","specification","of","the","clusters",".","See","above","for","more","information",".","Raises",":","ValueError",":","The","sum","of","the","token","counts","across","the","clusters","does","not","equal","the","vocabulary","size","or","a","length_projection_factor","is","not","in","the","range","(","0",",","1","]","."]
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,
      clusters: Sequence[Dict[str, Union[int, float]]] = gin.REQUIRED,
  ):
      """Configurable embedding for the vocabulary.
  
      Most of the arguments get passed to `mtf.layers.embedding_weights`.
  
      The clustering parameters are specified by the `clusters` argument. It is a
      list of dicts with keys:
        - token_count: The number of tokens in the cluster.
        - embedding_size: (optional) The hidden dimension size of the cluster's
          embedding. Defaults to the model dimension size.
        - length_projection_factor: (optional) Since MTF can't handle variable
          length dimensions, we project from the sequence length dimension to a
          dimension of size length_projection_factor * sequence_length during
          training. This can save compute time and resources if the cluster has
          many tokens that appear infrequently. If all of the tokens belonging to
          the cluster cannot fit within this reduced dimension, some will be
          discarded and ignored for the purposes of computing loss. Defaults 1.
          Ignored for the head (first) cluster and not during training.
  
      The first cluster will become the head cluster.
  
      For example, let's say we have a vocab size of 500k and pass as clusters:
        [
          {"token_count": 50000,  "embedding_size": 1024},
          {"token_count": 100000, "embedding_size": 256},
          {"token_count": 350000, "embedding_size": 64},
        ]
      Then tokens with ids 0 (inclusive) to 50k (exclusive) will be in the first
      cluster with embedding size of 1024, tokens with ids 50k to 150k will be in
      the second cluster with embedding size of 256, and tokens with ids 150k to
      500k will be in the third cluster with embedding size of 64.
  
      Args:
        mesh: a mtf.Mesh, the mesh used to layout the tensors.
        vocab_dim: a mtf.Dimension, the dimension corresponding to vocabulary.
        output_dim: a mtf.Dimension, the dimension corresponding to the model
          hidden states.
        variable_dtype: a mtf.VariableDType, the datatype information for the
          variables used in the embedding tensors.
        name: a string, a name to base variable names off of.
        ensemble_dim: a mtf.Dimension, the dimension used for ensembling.
          Absolutely no guarantees that this code will work with ensembling.
        clusters: a list(dict), specification of the clusters. See above for more
          information.
  
      Raises:
        ValueError: The sum of the token counts across the clusters does not equal
          the vocabulary size or a length_projection_factor is not in the range
          (0, 1].
      """
      self._mesh = mesh
      self._variable_dtype = variable_dtype
      self._name = name
      self._ensemble_dim = ensemble_dim
      self._vocab_dim = vocab_dim
      self._output_dim = output_dim
      self._num_clusters = len(clusters)
  
      token_counts = [cluster["token_count"] for cluster in clusters]
      if sum(token_counts) != vocab_dim.size:
          raise ValueError(
              "The cluster token counts {} do not sum to the vocab size {}.".format(
                  token_counts, vocab_dim.size
              )
          )
  
      self._tail_clusters = []
      start_token_id = 0
      for i, cluster_spec in enumerate(clusters):
          cluster = self._create_cluster(
              cluster_spec, i, start_token_id
          )
          if i == 0:
              self._head_cluster = cluster
          else:
              self._tail_clusters.append(cluster)
          start_token_id += cluster_spec["token_count"]
  
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 1 to create a new AdaptiveSoftmaxVocabEmbedding object from inherited base classes: object with example: obj = AdaptiveSoftmaxVocabEmbedding(mesh, vocab_dim, output_dim, variable_dtype, name, ensemble_dim, clusters)