def beam_search(symbols_to_logits_fn,
                initial_ids,
                beam_size,
                decode_length,
                vocab_size,
                alpha,
                eos_id,
                batch_size,
                minimum_score=None):
  """Beam search with length penalties.

  Requires a function that can take the currently decoded sybmols and return
  the logits for the next symbol. The implementation is inspired by
  https://arxiv.org/abs/1609.08144.

  Args:
    symbols_to_logits_fn: Interface to the model, to provide logits.
        Shoud take [batch_size, decoded_ids] and return [batch_size, vocab_size]
    initial_ids: Ids to start off the decoding, this will be the first thing
        handed to symbols_to_logits_fn (after expanding to beam size)
        [batch_size]
    beam_size: Size of the beam.
    decode_length: Number of steps to decode for.
    vocab_size: Size of the vocab, must equal the size of the logits returned by
        symbols_to_logits_fn
    alpha: alpha for length penalty.
    eos_id: ID for end of sentence.
    batch_size: Integer specifying batch size. Should match first dimension of
        initial_ids.
    minimum_score: Minimum score to continue exploring beam. Used to optimize
        performance. Ignored if None or 0. Should be in range [0-1].
  Returns:
    Tuple of
    (decoded beams [batch_size, beam_size, decode_length]
     decoding probablities [batch_size, beam_size])
  """

  def _inner_loop(i, alive_seq, alive_log_probs, finished_seq, finished_scores,
                  finished_flags):
    """Inner beam seach loop.

    There are three groups of tensors, alive, finished, and topk.
    The alive group contains information about the current alive sequences
    The topk group contains information about alive + topk current decoded words
    the finished group contains information about finished sentences, that is,
    the ones that have decoded to <EOS>. These are what we return.
    The general beam search algorithm is as follows:
    While we haven't terminated (pls look at termination condition)
      1. Grow the current alive to get beam*2 topk sequences
      2. Among the topk, keep the top beam_size ones that haven't reached EOS
      into alive
      3. Among the topk, keep the top beam_size ones have reached EOS into
      finished
    Repeat
    To make things simple with using fixed size tensors, we will end
    up inserting unfinished sequences into finished in the beginning. To stop
    that we add -ve INF to the score of the unfinished sequence so that when a
    true finished sequence does appear, it will have a higher score than all the
    unfinished ones.

    Args:
      i: loop index
      alive_seq: Topk sequences decoded so far.
          Shape is [batch_size, beam_size, decode_length + 1]
      alive_log_probs: probabilities of the beams. [batch_size, beam_size]
      finished_seq: Current finished sequences.
        [batch_size, beam_size, decode_length + 1]
      finished_scores: scores for each of these sequences.
        [batch_size, beam_size]
      finished_flags: finished bools for each of these sequences.
        [batch_size, beam_size]

    Returns:
      Tuple of
        (Incremented loop index
         New alive sequences,
         Log probs of the alive sequences,
         New finished sequences,
         Scores of the new finished sequences,
         Flags inidicating which sequence in finished as reached EOS)
    """

    # Each inner loop, we carry out three steps:
    # 1. Get the current topk items.
    # 2. Extract the ones that have finished and haven't finished
    # 3. Recompute the contents of finished based on scores.
    topk_seq, topk_log_probs, topk_scores, topk_finished = _grow_topk(
        i, alive_seq, alive_log_probs, batch_size, beam_size,
        symbols_to_logits_fn, alpha, vocab_size, eos_id, decode_length)
    alive_seq, alive_log_probs, _ = _grow_alive(topk_seq, topk_scores,
                                                topk_log_probs, topk_finished,
                                                batch_size, beam_size)
    finished_seq, finished_scores, finished_flags = _grow_finished(
        finished_seq, finished_scores, finished_flags, topk_seq, topk_scores,
        topk_finished, batch_size, beam_size)

    return (i + 1, alive_seq, alive_log_probs, finished_seq, finished_scores,
            finished_flags)

  def _loop_cond(i, unused_alive_seq, alive_log_probs, unused_finished_seq,
                 finished_scores, finished_in_finished):
    """Checking termination condition.

    We terminate when we decoded up to decode_length or the lowest scoring item
    in finished has a greater score that the higest prob item in alive divided
    by the max length penalty. Optionally also terminate if all alive scores
    are below lower bound.

    Args:
      i: loop index
      alive_log_probs: probabilities of the beams. [batch_size, beam_size]
      finished_scores: scores for each of these sequences.
        [batch_size, beam_size]
      finished_in_finished: finished bools for each of these sequences.
        [batch_size, beam_size]

    Returns:
      True to continue the loop, False to stop.
    """
    max_length_penalty = tf.pow(((5. + tf.to_float(decode_length)) / 6.), alpha)
    # The best possible score of the most likley alive sequence
    lower_bound_alive_scores = alive_log_probs[:, 0] / max_length_penalty

    # Now to compute the lowest score of a finished sequence in finished
    # If the sequence isn't finished, we multiply it's score by 0. since
    # scores are all -ve, taking the min will give us the score of the lowest
    # finished item.
    lowest_score_of_finished_in_finished = tf.reduce_min(
        finished_scores * tf.to_float(finished_in_finished), axis=1)
    # If none of the sequences have finished, then the min will be 0 and
    # we have to replace it by -ve INF if it is. The score of any seq in alive
    # will be much higher than -ve INF and the termination condition will not
    # be met.
    lowest_score_of_finished_in_finished = _apply_negative_infinity_mask(
        lowest_score_of_finished_in_finished,
        tf.logical_not(tf.reduce_any(finished_in_finished, 1)))

    # Will terminate beam search early if bound_is_met is True.
    bound_is_met = tf.reduce_all(
        tf.greater(lowest_score_of_finished_in_finished,
                   lower_bound_alive_scores))

    # Check if all alive scores are below minimum.
    if minimum_score:
      minimum_score_log = tf.log(minimum_score)
      bound_is_met = tf.logical_or(
          bound_is_met,
          tf.reduce_all(tf.less(lower_bound_alive_scores, minimum_score_log)))

    return tf.logical_and(
        tf.less(i, decode_length), tf.logical_not(bound_is_met))

  # Assume initial_ids are prob 1.0
  initial_log_probs = tf.constant([[0.] + [-float("inf")] * (beam_size - 1)])
  # Expand size to [batch_size, beam_size].
  alive_log_probs = tf.tile(initial_log_probs, [batch_size, 1])

  # Expand size to [batch_size, beam_size, decode_length + 1]
  alive_seq = tf.expand_dims(initial_ids, 1)  # [batch_size, 1]
  alive_seq = tf.tile(alive_seq, [1, beam_size])  # [batch_size, beam_size]
  alive_seq = _one_hot_tensor_3d(alive_seq, 0, decode_length + 1)

  # Finished will keep track of all the sequences that have finished so far
  # Finished log probs will be negative infinity in the beginning
  # finished_flags will keep track of booleans
  finished_seq = tf.zeros(tf.shape(alive_seq), tf.int32)
  # Setting the scores of the initial to negative infinity.
  finished_scores = tf.ones([batch_size, beam_size], dtype=tf.float32) * -INF
  finished_flags = tf.zeros([batch_size, beam_size], tf.bool)

  initial_variables = [
      tf.constant(0),  # []
      alive_seq,  # [batch_size, beam_size, decode_length + 1]
      alive_log_probs,  # [batch_size, beam_size]
      finished_seq,  # [batch_size, beam_size, decode_length + 1]
      finished_scores,  # [batch_size, beam_size]
      finished_flags  # [batch_size, beam_size]
  ]

  # Execute while loop.
  (_, alive_seq, alive_log_probs, finished_seq, finished_scores,
   finished_flags) = tf.while_loop(
       _loop_cond,
       _inner_loop,
       initial_variables,
       parallel_iterations=1,
       back_prop=False)

  # Accounting for corner case: It's possible that no sequence in alive for a
  # particular batch item ever reached EOS. In that case, we should just copy
  # the contents of alive for that batch item. tf.reduce_any(finished_flags, 1)
  # if 0, means that no sequence for that batch index had reached EOS. We need
  # to do the same for the scores as well.
  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)
  finished_seq = tf.cast(finished_seq, tf.int64)
  return finished_seq, finished_scores
Beam search with length penalties.

  Requires a function that can take the currently decoded sybmols and return
  the logits for the next symbol. The implementation is inspired by
  https://arxiv.org/abs/1609.08144.

  Args:
    symbols_to_logits_fn: Interface to the model, to provide logits.
        Shoud take [batch_size, decoded_ids] and return [batch_size, vocab_size]
    initial_ids: Ids to start off the decoding, this will be the first thing
        handed to symbols_to_logits_fn (after expanding to beam size)
        [batch_size]
    beam_size: Size of the beam.
    decode_length: Number of steps to decode for.
    vocab_size: Size of the vocab, must equal the size of the logits returned by
        symbols_to_logits_fn
    alpha: alpha for length penalty.
    eos_id: ID for end of sentence.
    batch_size: Integer specifying batch size. Should match first dimension of
        initial_ids.
    minimum_score: Minimum score to continue exploring beam. Used to optimize
        performance. Ignored if None or 0. Should be in range [0-1].
  Returns:
    Tuple of
    (decoded beams [batch_size, beam_size, decode_length]
     decoding probablities [batch_size, beam_size])
  
language/xsp/model/beam_search.py