Gather the top k sequences that are still alive.

    Args:
      new_seq: New sequences generated by growing the current alive sequences
        int32 tensor with shape [batch_size, 2 * beam_size, cur_index + 1]
      new_log_probs: Log probabilities of new sequences
        float32 tensor with shape [batch_size, beam_size]
      new_cache: Dict of cached values for each sequence.

    Returns:
      Dictionary with alive keys from _StateKeys:
        {Top beam_size sequences that are still alive (don't end with eos_id)
         Log probabilities of top alive sequences
         Dict cache storing decoder states for top alive sequences}

def _get_new_alive_state(self, new_seq, new_log_probs, new_cache):
    """Gather the top k sequences that are still alive.

    Args:
      new_seq: New sequences generated by growing the current alive sequences
        int32 tensor with shape [batch_size, 2 * beam_size, cur_index + 1]
      new_log_probs: Log probabilities of new sequences
        float32 tensor with shape [batch_size, beam_size]
      new_cache: Dict of cached values for each sequence.

    Returns:
      Dictionary with alive keys from _StateKeys:
        {Top beam_size sequences that are still alive (don't end with eos_id)
         Log probabilities of top alive sequences
         Dict cache storing decoder states for top alive sequences}
    """
    # To prevent finished sequences from being considered, set log probs to -INF
    new_finished_flags = tf.equal(new_seq[:, :, -1], self.eos_id)
    new_log_probs += tf.to_float(new_finished_flags) * -INF

    top_alive_seq, top_alive_log_probs, top_alive_cache = _gather_topk_beams(
        [new_seq, new_log_probs, new_cache], new_log_probs, self.batch_size,
        self.beam_size)

    return {
        _StateKeys.ALIVE_SEQ: top_alive_seq,
        _StateKeys.ALIVE_LOG_PROBS: top_alive_log_probs,
        _StateKeys.ALIVE_CACHE: top_alive_cache
    }