def _tokens_to_logits(last_token, cache, model, internal_state=None):
  """Computes the next token logits.

  Args:
    last_token: An array of shape (batch_size, 1) containing last token ids.
    cache: A flax.nn.attention.Cache object.
    model: A Jax decoder model to be used for computing the next token logits.
    internal_state: A dict with internal state received from the previous time
      step. If None, no information is shared across time steps.

  Returns:
    logits: An array of shape (batch_size, vocab_size) with the logits.
    new_cache: A flax.nn.attention.Cache object with the updated cache.
    new_internal_state: A dict with internal state passed to the next time step.
  """
  del internal_state  # Not used.
  # The returned logits have shape (batch_size, 1, vocab_size).
  with cache.mutate() as new_cache:
    logits = model(last_token, train=False, cache=new_cache)

  # Remove the singleton dimension to return shape (batch_size, vocab_size).
  logits = logits.squeeze(axis=1)
  return logits, new_cache, None