(symbols_to_logits_fn,
                init_seq_BxT,
                initial_cache_BxU,
                vocab_size,
                beam_size,
                length_norm_fn,
                eos_id=1)
return final_seq_BxMxT, final_scores_BxM
Beam search.

  Args:
    symbols_to_logits_fn: fn(seq_BxT, cache_BxU, i) -> (logits_BxV, cache_BxU)
    init_seq_BxT: initial sequence ids.
    initial_cache_BxU: dictionary of tensors with shape BxU.
    vocab_size: vocabulary size.
    beam_size: beam size.
    length_norm_fn: length normalization function.
    eos_id: end of sequence.

  Returns:
    Tuple of (beams_BxMxT, scores_BxM). Beam searched sequences and scores.
Beam search.
Beam search .
def beam_search(symbols_to_logits_fn,
                init_seq_BxT,
                initial_cache_BxU,
                vocab_size,
                beam_size,
                length_norm_fn,
                eos_id=1):
  """Beam search.

  Args:
    symbols_to_logits_fn: fn(seq_BxT, cache_BxU, i) -> (logits_BxV, cache_BxU)
    init_seq_BxT: initial sequence ids.
    initial_cache_BxU: dictionary of tensors with shape BxU.
    vocab_size: vocabulary size.
    beam_size: beam size.
    length_norm_fn: length normalization function.
    eos_id: end of sequence.

  Returns:
    Tuple of (beams_BxMxT, scores_BxM). Beam searched sequences and scores.
  """
  B, T = init_seq_BxT.shape
  M, V = beam_size, vocab_size
  dtype = tf.float32
  int_dtype = init_seq_BxT.dtype

  def _loop_body(i, alive_seq_BxMxT, alive_log_probs_BxM, alive_cache_BxMxU,
                 finished_seq_BxMxT, finished_scores_BxM):
    """Beam search loop body."""
    # Decode one step with beam
    logits_BMxV, cache_BMxU = symbols_to_logits_fn(
        _flatten_beam_dim(alive_seq_BxMxT),
        tf.nest.map_structure(_flatten_beam_dim, alive_cache_BxMxU), i)
    logits_BxMxV = _unflatten_beam_dim(logits_BMxV, M)
    new_cache_BxMxU = tf.nest.map_structure(lambda t: _unflatten_beam_dim(t, M),
                                            cache_BMxU)

    # select top 2 * beam_size and fill alive and finished.
    log_probs_BxMxV = logits_BxMxV - tf.reduce_logsumexp(
        logits_BxMxV, axis=2, keepdims=True)
    log_probs_BxMxV += tf.expand_dims(alive_log_probs_BxM, axis=2)
    log_probs_BxMV = tf.reshape(log_probs_BxMxV, [B, -1])
    new_log_probs_Bx2M, topk_indices_Bx2M = tf.nn.top_k(log_probs_BxMV, k=2 * M)
    topk_beam_Bx2M = topk_indices_Bx2M // V
    topk_seq_Bx2MxT, new_cache_Bx2MxU = _gather_nested(
        [alive_seq_BxMxT, new_cache_BxMxU], topk_beam_Bx2M)
    topk_ids_Bx2M = topk_indices_Bx2M % V
    new_seq_Bx2MxT = _update_i(topk_seq_Bx2MxT, topk_ids_Bx2M, i)
    new_finished_flags_Bx2M = tf.cast(
        tf.reduce_any(tf.equal(new_seq_Bx2MxT, eos_id), axis=-1), dtype)

    # get new alive
    _, topk_alive_indices_BxM = tf.nn.top_k(
        new_log_probs_Bx2M + new_finished_flags_Bx2M * dtype.min, k=M)
    (alive_seq_BxMxT, alive_log_probs_BxM, alive_cache_BxMxU) = _gather_nested(
        [new_seq_Bx2MxT, new_log_probs_Bx2M, new_cache_Bx2MxU],
        topk_alive_indices_BxM)

    # get new finished
    new_scores_Bx2M = length_norm_fn(new_log_probs_Bx2M, i + 1)
    new_scores_Bx2M += (1 - new_finished_flags_Bx2M) * dtype.min
    finished_seq_Bx3MxT = tf.concat([finished_seq_BxMxT, new_seq_Bx2MxT],
                                    axis=1)
    finished_scores_Bx3M = tf.concat([finished_scores_BxM, new_scores_Bx2M],
                                     axis=1)
    _, topk_finished_indices_BxM = tf.nn.top_k(finished_scores_Bx3M, k=M)
    (finished_seq_BxMxT, finished_scores_BxM) = _gather_nested(
        [finished_seq_Bx3MxT, finished_scores_Bx3M], topk_finished_indices_BxM)

    return [
        i + 1, alive_seq_BxMxT, alive_log_probs_BxM, alive_cache_BxMxU,
        finished_seq_BxMxT, finished_scores_BxM
    ]

  # initialize.
  init_i = tf.constant(0, dtype=int_dtype)
  init_alive_seq_BxMxT = _expand_to_beam_size(init_seq_BxT, M)
  log_probs_1xM = tf.constant([[0.] + [dtype.min] * (M - 1)], dtype=dtype)
  init_alive_log_probs_BxM = tf.tile(log_probs_1xM, [B, 1])
  init_alive_cache_BxMxU = tf.nest.map_structure(
      lambda t: _expand_to_beam_size(t, M), initial_cache_BxU)
  init_finished_seq_BxMxT = tf.zeros(tf.shape(init_alive_seq_BxMxT), int_dtype)
  init_finished_scores_BxM = tf.zeros([B, M], dtype=dtype) + dtype.min

  # run loop.
  (_, final_alive_seq_BxMxT, final_alive_scores_BxM, _,
   final_finished_seq_BxMxT, final_finished_scores_BxM) = tf.while_loop(
       lambda *args: True,  # Always do T iterations
       _loop_body,
       loop_vars=[
           init_i, init_alive_seq_BxMxT, init_alive_log_probs_BxM,
           init_alive_cache_BxMxU, init_finished_seq_BxMxT,
           init_finished_scores_BxM
       ],
       parallel_iterations=1,
       back_prop=False,
       maximum_iterations=T,
   )

  # process finished.
  final_finished_flag_BxMx1 = tf.reduce_any(
      tf.equal(final_finished_seq_BxMxT, eos_id), axis=-1, keepdims=True)
  final_seq_BxMxT = tf.where(
      tf.tile(final_finished_flag_BxMx1, [1, 1, T]), final_finished_seq_BxMxT,
      final_alive_seq_BxMxT)
  final_scores_BxM = tf.where(
      tf.squeeze(final_finished_flag_BxMx1, axis=-1), final_finished_scores_BxM,
      final_alive_scores_BxM)
  return final_seq_BxMxT, final_scores_BxM
def beam_search ( symbols_to_logits_fn , init_seq_BxT , initial_cache_BxU , vocab_size , beam_size , length_norm_fn , eos_id = 1 ) : B , T = init_seq_BxT . shape M , V = beam_size , vocab_size dtype = tf . float32 int_dtype = init_seq_BxT . dtype def _loop_body ( i , alive_seq_BxMxT , alive_log_probs_BxM , alive_cache_BxMxU , finished_seq_BxMxT , finished_scores_BxM ) : """Beam search loop body.""" # Decode one step with beam logits_BMxV , cache_BMxU = symbols_to_logits_fn ( _flatten_beam_dim ( alive_seq_BxMxT ) , tf . nest . map_structure ( _flatten_beam_dim , alive_cache_BxMxU ) , i ) logits_BxMxV = _unflatten_beam_dim ( logits_BMxV , M ) new_cache_BxMxU = tf . nest . map_structure ( lambda t : _unflatten_beam_dim ( t , M ) , cache_BMxU ) # select top 2 * beam_size and fill alive and finished. log_probs_BxMxV = logits_BxMxV - tf . reduce_logsumexp ( logits_BxMxV , axis = 2 , keepdims = True ) log_probs_BxMxV += tf . expand_dims ( alive_log_probs_BxM , axis = 2 ) log_probs_BxMV = tf . reshape ( log_probs_BxMxV , [ B , - 1 ] ) new_log_probs_Bx2M , topk_indices_Bx2M = tf . nn . top_k ( log_probs_BxMV , k = 2 * M ) topk_beam_Bx2M = topk_indices_Bx2M // V topk_seq_Bx2MxT , new_cache_Bx2MxU = _gather_nested ( [ alive_seq_BxMxT , new_cache_BxMxU ] , topk_beam_Bx2M ) topk_ids_Bx2M = topk_indices_Bx2M % V new_seq_Bx2MxT = _update_i ( topk_seq_Bx2MxT , topk_ids_Bx2M , i ) new_finished_flags_Bx2M = tf . cast ( tf . reduce_any ( tf . equal ( new_seq_Bx2MxT , eos_id ) , axis = - 1 ) , dtype ) # get new alive _ , topk_alive_indices_BxM = tf . nn . top_k ( new_log_probs_Bx2M + new_finished_flags_Bx2M * dtype . min , k = M ) ( alive_seq_BxMxT , alive_log_probs_BxM , alive_cache_BxMxU ) = _gather_nested ( [ new_seq_Bx2MxT , new_log_probs_Bx2M , new_cache_Bx2MxU ] , topk_alive_indices_BxM ) # get new finished new_scores_Bx2M = length_norm_fn ( new_log_probs_Bx2M , i + 1 ) new_scores_Bx2M += ( 1 - new_finished_flags_Bx2M ) * dtype . min finished_seq_Bx3MxT = tf . concat ( [ finished_seq_BxMxT , new_seq_Bx2MxT ] , axis = 1 ) finished_scores_Bx3M = tf . concat ( [ finished_scores_BxM , new_scores_Bx2M ] , axis = 1 ) _ , topk_finished_indices_BxM = tf . nn . top_k ( finished_scores_Bx3M , k = M ) ( finished_seq_BxMxT , finished_scores_BxM ) = _gather_nested ( [ finished_seq_Bx3MxT , finished_scores_Bx3M ] , topk_finished_indices_BxM ) return [ i + 1 , alive_seq_BxMxT , alive_log_probs_BxM , alive_cache_BxMxU , finished_seq_BxMxT , finished_scores_BxM ] # initialize. init_i = tf . constant ( 0 , dtype = int_dtype ) init_alive_seq_BxMxT = _expand_to_beam_size ( init_seq_BxT , M ) log_probs_1xM = tf . constant ( [ [ 0. ] + [ dtype . min ] * ( M - 1 ) ] , dtype = dtype ) init_alive_log_probs_BxM = tf . tile ( log_probs_1xM , [ B , 1 ] ) init_alive_cache_BxMxU = tf . nest . map_structure ( lambda t : _expand_to_beam_size ( t , M ) , initial_cache_BxU ) init_finished_seq_BxMxT = tf . zeros ( tf . shape ( init_alive_seq_BxMxT ) , int_dtype ) init_finished_scores_BxM = tf . zeros ( [ B , M ] , dtype = dtype ) + dtype . min # run loop. ( _ , final_alive_seq_BxMxT , final_alive_scores_BxM , _ , final_finished_seq_BxMxT , final_finished_scores_BxM ) = tf . while_loop ( lambda * args : True , # Always do T iterations _loop_body , loop_vars = [ init_i , init_alive_seq_BxMxT , init_alive_log_probs_BxM , init_alive_cache_BxMxU , init_finished_seq_BxMxT , init_finished_scores_BxM ] , parallel_iterations = 1 , back_prop = False , maximum_iterations = T , ) # process finished. final_finished_flag_BxMx1 = tf . reduce_any ( tf . equal ( final_finished_seq_BxMxT , eos_id ) , axis = - 1 , keepdims = True ) final_seq_BxMxT = tf . where ( tf . tile ( final_finished_flag_BxMx1 , [ 1 , 1 , T ] ) , final_finished_seq_BxMxT , final_alive_seq_BxMxT ) final_scores_BxM = tf . where ( tf . squeeze ( final_finished_flag_BxMx1 , axis = - 1 ) , final_finished_scores_BxM , final_alive_scores_BxM ) return final_seq_BxMxT , final_scores_BxM