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)

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, 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