def __call__(
      self, token_index: named_axes.NamedArray
  ) -> named_axes.NamedArray:
    """Retrieves tokens from the embedding table.

    Args:
      token_index: A named array of token indices. Axis names of this array must
        be disjoint from the axis names in the embedding table.

    Returns:
      A named array of embeddings, which includes all named axes of the input
      along with the (non-vocabulary) named axes of the embedding table.
    """
    return self.table.embeddings.value[
        {self.table.vocabulary_axis: token_index}
    ]
Retrieves tokens from the embedding table.

    Args:
      token_index: A named array of token indices. Axis names of this array must
        be disjoint from the axis names in the embedding table.

    Returns:
      A named array of embeddings, which includes all named axes of the input
      along with the (non-vocabulary) named axes of the embedding table.
    
penzai/deprecated/v1/nn/embeddings.py