class_docstr: Helper class for adaptive embeddings specifying a cluster of tokens.
  
  Essentially a wrapper around a vocab embedding for the cluster with additional
  metadata so that we can apply the embedding to the actual ids and hidden
  states.
class_docstr_tok: ["Helper","class","for","adaptive","embeddings","specifying","a","cluster","of","tokens",".","Essentially","a","wrapper","around","a","vocab","embedding","for","the","cluster","with","additional","metadata","so","that","we","can","apply","the","embedding","to","the","actual","ids","and","hidden","states","."]
docstr: Returns first end_index entries in log_softmax along the vocab dim.
docstr_tok: ["Returns","first","end_index","entries","in","log_softmax","along","the","vocab","dim","."]
code: def get_log_softmax_prefix(self, log_softmax, end_index):
      """Returns first end_index entries in log_softmax along the vocab dim."""
      prefix_dim = mtf.Dimension(self._vocab_dim.name, end_index)
  
      indices = mtf.mtf_range(
          log_softmax.mesh, dim=self._vocab_dim, dtype=tf.int32
      )
      prefix_indices = mtf.where(
          mtf.less(indices, end_index), indices, -1
      )
      projection = mtf.one_hot(
          prefix_indices, prefix_dim, dtype=log_softmax.dtype
      )
  
      return mtf.einsum(
          [log_softmax, projection], reduced_dims=[self._vocab_dim]
      )
  
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 10 for calling the _Cluster obj's underlying member method code with example usage: obj.get_log_softmax_prefix(log_softmax, end_index) and returns: mtf