docs: def search(self, initial_ids, initial_cache):
      """Beam search for sequences with highest scores."""
      state, state_shapes = self._create_initial_state(initial_ids, initial_cache)
  
      finished_state = tf.while_loop(
          self._continue_search, self._search_step, loop_vars=[state],
          shape_invariants=[state_shapes], parallel_iterations=1, back_prop=False)
      finished_state = finished_state[0]
  
      alive_seq = finished_state[_StateKeys.ALIVE_SEQ]
      alive_log_probs = finished_state[_StateKeys.ALIVE_LOG_PROBS]
      finished_seq = finished_state[_StateKeys.FINISHED_SEQ]
      finished_scores = finished_state[_StateKeys.FINISHED_SCORES]
      finished_flags = finished_state[_StateKeys.FINISHED_FLAGS]
  
      # Account for corner case where there are no finished sequences for a
      # particular batch item. In that case, return alive sequences for that batch
      # item.
      finished_seq = tf.where(
          tf.reduce_any(finished_flags, 1), finished_seq, alive_seq)
      finished_scores = tf.where(
          tf.reduce_any(finished_flags, 1), finished_scores, alive_log_probs)
      return finished_seq, finished_scores