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