code: def compute_topk_scores_and_seq(
      sequences,
      scores,
      scores_to_gather,
      flags,
      beam_dim,
      prefix="default",
  ):
      """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 two 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.
  
      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_dim: mtf.Dimension
        prefix: an optional string
      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],
         selector)
      """
      (
          unused_batch_dim,
          old_beam_dim,
          unused_length_dim,
      ) = sequences.shape.dims
      _, topk_indices = mtf.top_k(scores, old_beam_dim, k_dim=beam_dim)
  
      selector = mtf.one_hot(
          topk_indices, old_beam_dim, dtype=tf.float32
      )
  
      # 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):
          with tf.name_scope(prefix + name):
              output_shape = mtf.Shape(
                  [
                      beam_dim if d == old_beam_dim else d
                      for d in tensor.shape.dims
                  ]
              )
              return mtf.gather(
                  tensor,
                  topk_indices,
                  old_beam_dim,
                  output_shape=output_shape,
              )
  
      topk_seq = gather(sequences, "_seq")
      topk_flags = gather(flags, "_flags")
      topk_gathered_scores = gather(scores_to_gather, "_scores")
      return topk_seq, topk_gathered_scores, topk_flags, selector
  
file_path: mesh/mesh_tensorflow/beam_search.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import gin
  from mesh_tensorflow import ops_with_redefined_builtins
  import tensorflow.compat.v1
label_desc: Use image node_id 1 for calling a global function with example usage: compute_topk_scores_and_seq(sequences, scores, scores_to_gather, flags, beam_dim, prefix) and returns: topk_seq,  topk_gathered_scores,  topk_flags,  selector, mtf