def _grow_topk(i, alive_seq, alive_log_probs, batch_size, beam_size,
               symbols_to_logits_fn, alpha, vocab_size, eos_id, decode_length):
  r"""Inner beam seach loop.

  This function takes the current alive sequences, and grows them to topk
  sequences where k = 2*beam. We use 2*beam because, we could have beam_size
  number of sequences that might hit <EOS> and there will be no alive
  sequences to continue. With 2*beam_size, this will not happen. This relies
  on the assumption the vocab size is > beam size. If this is true, we'll
  have at least beam_size non <EOS> extensions if we extract the next top
  2*beam words.
  Length penalty is given by = (5+len(decode)/6) ^ -\alpha. Pls refer to
  https://arxiv.org/abs/1609.08144.

  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 these sequences. [batch_size, beam_size]
    batch_size: Integer specifying batch size.
    beam_size: Integer specifying beam size.
    symbols_to_logits_fn: Interface to the model, to provide logits.
        Shoud take [batch_size, decoded_ids] and return [batch_size, vocab_size]
    alpha: alpha for length penalty.
    vocab_size: Size of the vocab, must equal the size of the logits returned by
        symbols_to_logits_fn
    eos_id: ID for end of sentence.
    decode_length: Maximum length for decoded sequence.
  Returns:
    Tuple of
      (Topk sequences extended by the next word,
       The log probs of these sequences,
       The scores with length penalty of these sequences,
       Flags indicating which of these sequences have finished decoding)
  """
  # Get the logits for all the possible next symbols
  flat_ids = tf.reshape(alive_seq, [batch_size * beam_size, -1])

  # (batch_size * beam_size, decoded_length)
  flat_logits = symbols_to_logits_fn(i, flat_ids)
  logits = tf.reshape(flat_logits, (batch_size, beam_size, -1))

  # Convert logits to normalized log probs
  candidate_log_probs = _log_prob_from_logits(logits)

  # Multiply the probabilites by the current probabilites of the beam.
  # (batch_size, beam_size, vocab_size) + (batch_size, beam_size, 1)
  log_probs = candidate_log_probs + tf.expand_dims(alive_log_probs, axis=2)

  length_penalty = tf.pow(((5. + tf.to_float(i + 1)) / 6.), alpha)

  # curr_scores has shape [batch_size, beam_size, vocab_size].
  curr_scores = log_probs / length_penalty
  # Flatten out to have shape [batch_size, beam_size * vocab_size].
  # Note that vocab size is not always known statically because the extended
  # vocab size can vary based on the input length.
  flat_curr_scores = tf.reshape(curr_scores, [batch_size, -1])

  topk_scores, topk_ids = tf.nn.top_k(flat_curr_scores, beam_size * 2)

  # Recovering the log probs because we will need to send them back
  topk_log_probs = topk_scores * length_penalty

  # Work out what beam the top probs are in.
  topk_beam_index = topk_ids // vocab_size
  topk_ids %= vocab_size  # Unflatten the ids

  # The next three steps are to create coordinates for tf.gather_nd to pull
  # out the correct seqences from id's that we need to grow.
  # We will also use the coordinates to gather the booleans of the beam items
  # that survived.
  batch_pos = _compute_batch_indices(batch_size, beam_size * 2)

  # top beams 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.
  topk_coordinates = tf.stack([batch_pos, topk_beam_index], axis=2)

  # Gather up the most probable 2*beams both for the ids and finished_in_alive
  # bools
  topk_seq = tf.gather_nd(alive_seq, topk_coordinates)

  # Append the most probable alive
  topk_ids_padded = _one_hot_tensor_3d(topk_ids, i + 1, decode_length + 1)
  topk_seq += topk_ids_padded

  topk_finished = tf.equal(topk_ids, eos_id)

  return topk_seq, topk_log_probs, topk_scores, topk_finished
Inner beam seach loop.

  This function takes the current alive sequences, and grows them to topk
  sequences where k = 2*beam. We use 2*beam because, we could have beam_size
  number of sequences that might hit <EOS> and there will be no alive
  sequences to continue. With 2*beam_size, this will not happen. This relies
  on the assumption the vocab size is > beam size. If this is true, we'll
  have at least beam_size non <EOS> extensions if we extract the next top
  2*beam words.
  Length penalty is given by = (5+len(decode)/6) ^ -\alpha. Pls refer to
  https://arxiv.org/abs/1609.08144.

  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 these sequences. [batch_size, beam_size]
    batch_size: Integer specifying batch size.
    beam_size: Integer specifying beam size.
    symbols_to_logits_fn: Interface to the model, to provide logits.
        Shoud take [batch_size, decoded_ids] and return [batch_size, vocab_size]
    alpha: alpha for length penalty.
    vocab_size: Size of the vocab, must equal the size of the logits returned by
        symbols_to_logits_fn
    eos_id: ID for end of sentence.
    decode_length: Maximum length for decoded sequence.
  Returns:
    Tuple of
      (Topk sequences extended by the next word,
       The log probs of these sequences,
       The scores with length penalty of these sequences,
       Flags indicating which of these sequences have finished decoding)
  
language/xsp/model/beam_search.py