def tokens_ids_to_logits(flat_ids, flat_cache):
    """Token slice to logits from decoder model."""
    # --> [batch * beam, 1, vocab]
    flat_logits, new_vars = model.apply(
        {
            'params': state.optimizer.target,
            'cache': flat_cache
        },
        flat_ids,
        encoded_inputs,
        flat_ids > 0,
        encoded_mask,
        mutable=['cache'],
        method=models.Model.decode)
    new_flat_cache = new_vars['cache']
    # Remove singleton sequence-length dimension:
    # [batch * beam, 1, vocab] --> [batch * beam, vocab]
    flat_logits = flat_logits.squeeze(axis=1)
    return flat_logits, new_flat_cache
Token slice to logits from decoder model.
language/gscan/xattn_model/predict.py