Grow alive sequences by one token, and collect top 2*beam_size sequences.

    2*beam_size sequences are collected because some sequences may have reached
    the EOS token. 2*beam_size ensures that at least beam_size sequences are
    still alive.

    Args:
      state: A dictionary with the current loop state.
    Returns:
      Tuple of
      (Top 2*beam_size sequences [batch_size, 2 * beam_size, cur_index + 1],
       Scores of returned sequences [batch_size, 2 * beam_size],
       New alive cache, for each of the 2 * beam_size sequences)

def _grow_alive_seq(self, state):
    """Grow alive sequences by one token, and collect top 2*beam_size sequences.

    2*beam_size sequences are collected because some sequences may have reached
    the EOS token. 2*beam_size ensures that at least beam_size sequences are
    still alive.

    Args:
      state: A dictionary with the current loop state.
    Returns:
      Tuple of
      (Top 2*beam_size sequences [batch_size, 2 * beam_size, cur_index + 1],
       Scores of returned sequences [batch_size, 2 * beam_size],
       New alive cache, for each of the 2 * beam_size sequences)
    """
    i = state[_StateKeys.CUR_INDEX]
    alive_seq = state[_StateKeys.ALIVE_SEQ]
    alive_log_probs = state[_StateKeys.ALIVE_LOG_PROBS]
    alive_cache = state[_StateKeys.ALIVE_CACHE]

    beams_to_keep = 2 * self.beam_size

    # Get logits for the next candidate IDs for the alive sequences. Get the new
    # cache values at the same time.
    flat_ids = _flatten_beam_dim(alive_seq)  # [batch_size * beam_size]
    flat_cache = nest.map_structure(_flatten_beam_dim, alive_cache)

    flat_logits, flat_cache = self.symbols_to_logits_fn(flat_ids, i, flat_cache)

    # Unflatten logits to shape [batch_size, beam_size, vocab_size]
    logits = _unflatten_beam_dim(flat_logits, self.batch_size, self.beam_size)
    new_cache = nest.map_structure(
        lambda t: _unflatten_beam_dim(t, self.batch_size, self.beam_size),
        flat_cache)

    # Convert logits to normalized log probs
    candidate_log_probs = _log_prob_from_logits(logits)

    # Calculate new log probabilities if each of the alive sequences were
    # extended # by the the candidate IDs.
    # Shape [batch_size, beam_size, vocab_size]
    log_probs = candidate_log_probs + tf.expand_dims(alive_log_probs, axis=2)

    # Each batch item has beam_size * vocab_size candidate sequences. For each
    # batch item, get the k candidates with the highest log probabilities.
    flat_log_probs = tf.reshape(log_probs,
                                [-1, self.beam_size * self.vocab_size])
    topk_log_probs, topk_indices = tf.nn.top_k(flat_log_probs, k=beams_to_keep)

    # Extract the alive sequences that generate the highest log probabilities
    # after being extended.
    topk_beam_indices = topk_indices // self.vocab_size
    topk_seq, new_cache = _gather_beams(
        [alive_seq, new_cache], topk_beam_indices, self.batch_size,
        beams_to_keep)

    # Append the most probable IDs to the topk sequences
    topk_ids = topk_indices % self.vocab_size
    topk_ids = tf.expand_dims(topk_ids, axis=2)
    topk_seq = tf.concat([topk_seq, topk_ids], axis=2)
    return topk_seq, topk_log_probs, new_cache