code: def compute_topk_scores_and_seq(
      sequences,
      scores,
      scores_to_gather,
      flags,
      beam_size,
      batch_size,
      prefix="default",
      states_to_gather=None,
  ):
      """Given sequences and scores, will gather the top k=beam size sequences."""
      _, topk_indexes = tf.nn.top_k(scores, k=beam_size)
      batch_pos = compute_batch_indices(batch_size, beam_size)
      top_coordinates = tf.stack([batch_pos, topk_indexes], axis=2)
  
      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: delta/delta/layers/utils_tf.py
imports: import math
  import numpy
  import tensorflow
  from tensorflow.python.util import nest
label_desc: Use image node_id 11 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) and returns: topk_seq,  topk_gathered_scores,  topk_flags,  topk_gathered_states, tf