docs: 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]