code: def compute_topk_scores_and_seq(
      sequences,
      scores,
      scores_to_gather,
      flags,
      beam_size,
      batch_size,
      prefix="default",
      states_to_gather=None,
      use_top_k_with_unique=True,
  ):
      """Given sequences and scores, will gather the top k=beam size sequences.
  
      This function is used to grow alive, and finished. It takes sequences,
      scores, and flags, and returns the top k from sequences, scores_to_gather,
      and flags based on the values in scores.
  
      This method permits easy introspection using tfdbg.  It adds three named ops
      that are prefixed by `prefix`:
        - _topk_seq: the tensor for topk_seq returned by this method.
        - _topk_flags: the tensor for topk_finished_flags returned by this method.
        - _topk_scores: the tensor for tokp_gathered_scores returned by this method.
  
      Args:
        sequences: Tensor of sequences that we need to gather from.
          [batch_size, beam_size, seq_length]
        scores: Tensor of scores for each sequence in sequences.
          [batch_size, beam_size]. We will use these to compute the topk.
        scores_to_gather: Tensor of scores for each sequence in sequences.
          [batch_size, beam_size]. We will return the gathered scores from here.
          Scores to gather is different from scores because for grow_alive, we will
          need to return log_probs, while for grow_finished, we will need to return
          the length penalized scores.
        flags: Tensor of bools for sequences that say whether a sequence has reached
          EOS or not
        beam_size: int
        batch_size: int
        prefix: string that will prefix unique names for the ops run.
        states_to_gather: dict (possibly nested) of decoding states.
        use_tpu: A bool, whether to compute topk scores and sequences on TPU.
        use_top_k_with_unique: bool, whether to use a fast (but decreased precision)
          top_k during TPU beam search.
  
      Returns:
        Tuple of
        (topk_seq [batch_size, beam_size, decode_length],
         topk_gathered_scores [batch_size, beam_size],
         topk_finished_flags[batch_size, beam_size])
      """
      _, topk_indexes = tf.nn.top_k(scores, k=beam_size)
      # The next three steps are to create coordinates for tf.gather_nd to pull
      # out the topk sequences from sequences based on scores.
      # batch pos is a tensor like [[0,0,0,0,],[1,1,1,1],..]. It says which
      # batch the beam item is in. This will create the i of the i,j coordinate
      # needed for the gather
      batch_pos = compute_batch_indices(batch_size, beam_size)
  
      # top coordinates will give us the actual coordinates to do the gather.
      # stacking will create a tensor of dimension batch * beam * 2, where the
      # last dimension contains the i,j gathering coordinates.
      top_coordinates = tf.stack([batch_pos, topk_indexes], axis=2)
  
      # Gather up the highest scoring sequences.  For each operation added, give
      # it a concrete name to simplify observing these operations with tfdbg.
      # Clients can capture these tensors by watching these node names.
      def gather(tensor, name):
          return tf.gather_nd(
              tensor, top_coordinates, name=(prefix + name)
          )
  
      topk_seq = gather(sequences, "_topk_seq")
      topk_flags = gather(flags, "_topk_flags")
      topk_gathered_scores = gather(scores_to_gather, "_topk_scores")
      if states_to_gather:
          topk_gathered_states = nest.map_structure(
              lambda state: gather(state, "_topk_states"),
              states_to_gather,
          )
      else:
          topk_gathered_states = states_to_gather
  
      return (
          topk_seq,
          topk_gathered_scores,
          topk_flags,
          topk_gathered_states,
      )
  
file_path: finetune/finetune/util/beam_search.py
imports: import math
  import numpy
  from finetune.util.shapes import shape_list
  import tensorflow
  from tensorflow.python.ops import inplace_ops
  from tensorflow.python.util import nest
label_desc: Use image node_id 12 for calling a global function with example usage: compute_topk_scores_and_seq(sequences, scores, scores_to_gather, flags, beam_size, batch_size, prefix, states_to_gather, use_top_k_with_unique) and returns: topk_seq,  topk_gathered_scores,  topk_flags,  topk_gathered_states, tf