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

    Args:
      out_embeddings: The output embeddings that should be mapped to token
        logits. Should be a named array that includes the same axes as the
        embedding table, except for the vocabulary axis, and may also include
        additional batch axes.

    Returns:
      A named array of logits, which includes all batch axes of the input
      along with the vocabulary axis of the embedding table.
    """
    contracting_axes = [
        name
        for name in self.table.embeddings.value.named_shape.keys()
        if name != self.table.vocabulary_axis
    ]
    return linear_and_affine.contract(
        contracting_axes, out_embeddings, self.table.embeddings.value
    )
Retrieves tokens from the embedding table.

    Args:
      out_embeddings: The output embeddings that should be mapped to token
        logits. Should be a named array that includes the same axes as the
        embedding table, except for the vocabulary axis, and may also include
        additional batch axes.

    Returns:
      A named array of logits, which includes all batch axes of the input
      along with the vocabulary axis of the embedding table.
    
penzai/deprecated/v1/nn/embeddings.py