docs: def _get_new_finished_state(self, state, new_seq, new_log_probs):
      """Combine new and old finished sequences, and gather the top k sequences.
  
      Args:
        state: A dictionary with the current loop state.
        new_seq: New sequences generated by growing the current alive sequences
          int32 tensor with shape [batch_size, beam_size, i + 1]
        new_log_probs: Log probabilities of new sequences
          float32 tensor with shape [batch_size, beam_size]
  
      Returns:
        Dictionary with finished keys from _StateKeys:
          {Top beam_size finished sequences based on score,
           Scores of finished sequences,
           Finished flags of finished sequences}
      """
      i = state[_StateKeys.CUR_INDEX]
      finished_seq = state[_StateKeys.FINISHED_SEQ]
      finished_scores = state[_StateKeys.FINISHED_SCORES]
      finished_flags = state[_StateKeys.FINISHED_FLAGS]
  
      # First append a column of 0-ids to finished_seq to increment the length.
      # New shape of finished_seq: [batch_size, beam_size, i + 1]
      finished_seq = tf.concat(
          [finished_seq,
           tf.zeros([self.batch_size, self.beam_size, 1], tf.int32)], axis=2)
  
      # Calculate new seq scores from log probabilities.
      length_norm = _length_normalization(self.alpha, i + 1)
      new_scores = new_log_probs / length_norm
  
      # Set the scores of the still-alive seq in new_seq to large negative values.
      new_finished_flags = tf.equal(new_seq[:, :, -1], self.eos_id)
      new_scores += (1. - tf.to_float(new_finished_flags)) * -INF
  
      # Combine sequences, scores, and flags.
      finished_seq = tf.concat([finished_seq, new_seq], axis=1)
      finished_scores = tf.concat([finished_scores, new_scores], axis=1)
      finished_flags = tf.concat([finished_flags, new_finished_flags], axis=1)
  
      # Return the finished sequences with the best scores.
      top_finished_seq, top_finished_scores, top_finished_flags = (
          _gather_topk_beams([finished_seq, finished_scores, finished_flags],
                             finished_scores, self.batch_size, self.beam_size))
  
      return {
          _StateKeys.FINISHED_SEQ: top_finished_seq,
          _StateKeys.FINISHED_SCORES: top_finished_scores,
          _StateKeys.FINISHED_FLAGS: top_finished_flags
      }