Beam search loop body.

    Grow alive sequences by a single ID. Sequences that have reached the EOS
    token are marked as finished. The alive and finished sequences with the
    highest log probabilities and scores are returned.

    A sequence's finished score is calculating by dividing the log probability
    by the length normalization factor. Without length normalization, the
    search is more likely to return shorter sequences.

    Args:
      state: A dictionary with the current loop state.

    Returns:
      new state dictionary.

def _search_step(self, state):
    """Beam search loop body.

    Grow alive sequences by a single ID. Sequences that have reached the EOS
    token are marked as finished. The alive and finished sequences with the
    highest log probabilities and scores are returned.

    A sequence's finished score is calculating by dividing the log probability
    by the length normalization factor. Without length normalization, the
    search is more likely to return shorter sequences.

    Args:
      state: A dictionary with the current loop state.

    Returns:
      new state dictionary.
    """
    # Grow alive sequences by one token.
    new_seq, new_log_probs, new_cache = self._grow_alive_seq(state)
    # Collect top beam_size alive sequences
    alive_state = self._get_new_alive_state(new_seq, new_log_probs, new_cache)

    # Combine newly finished sequences with existing finished sequences, and
    # collect the top k scoring sequences.
    finished_state = self._get_new_finished_state(state, new_seq, new_log_probs)

    # Increment loop index and create new state dictionary
    new_state = {_StateKeys.CUR_INDEX: state[_StateKeys.CUR_INDEX] + 1}
    new_state.update(alive_state)
    new_state.update(finished_state)
    return [new_state]