class_docstr: Transformer Encoder is stacked with several
  TransformerdecLayers, consist of beamsearch for infrence.
class_docstr_tok: ["Transformer","Encoder","is","stacked","with","several","TransformerdecLayers",",","consist","of","beamsearch","for","infrence","."]
docstr: Beam search with length penalties.
docstr_tok: ["Beam","search","with","length","penalties","."]
code: def beam_search(
      symbols_to_logits_fn,
      initial_ids,
      beam_size,
      decode_length,
      vocab_size,
      alpha,
      eos_id,
      states=None,
      stop_early=True,
      INF=1.0 * 1e20,
  ):
      """Beam search with length penalties."""
      batch_size = utils.shape_list(initial_ids)[0]
  
      initial_log_probs = tf.constant(
          [[0.0] + [-INF] * (beam_size - 1)]
      )
      # (batch_size, beam_size)
      alive_log_probs = tf.tile(initial_log_probs, [batch_size, 1])
  
      alive_seq = utils.expand_to_beam_size(initial_ids, beam_size)
      # (batch_size, beam_size, 1)
      alive_seq = tf.expand_dims(alive_seq, axis=2)
      if states:
          states = nest.map_structure(
              lambda state: utils.expand_to_beam_size(state, beam_size),
              states,
          )
      else:
          states = {}
  
      # (batch_size, beam_size, 1)
      finished_seq = tf.zeros(utils.shape_list(alive_seq), tf.int32)
      # (batch_size, beam_size)
      finished_scores = tf.ones([batch_size, beam_size]) * -INF
      # (batch_size, beam_size)
      finished_flags = tf.zeros([batch_size, beam_size], tf.bool)
  
      def grow_finished(
          finished_seq,
          finished_scores,
          finished_flags,
          curr_seq,
          curr_scores,
          curr_finished,
      ):
          """
          Given sequences and scores from finished sequence and current finished sequence
          , will gather the top k=beam size sequences to update finished seq.
          """
          # padding zero for finished seq
          finished_seq = tf.concat(
              [
                  finished_seq,
                  tf.zeros([batch_size, beam_size, 1], tf.int32),
              ],
              axis=2,
          )
  
          # mask unfinished curr seq
          curr_scores += (1.0 - tf.to_float(curr_finished)) * -INF
  
          # concatenating the sequences and scores along beam axis
          # (batch_size, 2xbeam_size, seq_len)
          curr_finished_seq = tf.concat(
              [finished_seq, curr_seq], axis=1
          )
          curr_finished_scores = tf.concat(
              [finished_scores, curr_scores], axis=1
          )
          curr_finished_flags = tf.concat(
              [finished_flags, curr_finished], axis=1
          )
          return utils.compute_topk_scores_and_seq(
              curr_finished_seq,
              curr_finished_scores,
              curr_finished_scores,
              curr_finished_flags,
              beam_size,
              batch_size,
              "grow_finished",
          )
  
      def grow_alive(
          curr_seq, curr_scores, curr_log_probs, curr_finished, states
      ):
          """Given sequences and scores, will gather the top k=beam size sequences."""
          curr_scores += tf.to_float(curr_finished) * -INF
          return utils.compute_topk_scores_and_seq(
              curr_seq,
              curr_scores,
              curr_log_probs,
              curr_finished,
              beam_size,
              batch_size,
              "grow_alive",
              states,
          )
  
      def grow_topk(i, alive_seq, alive_log_probs, states):
          """Inner beam search loop."""
          flat_ids = tf.reshape(alive_seq, [batch_size * beam_size, -1])
  
          # (batch_size * beam_size, decoded_length)
          if states:
              flat_states = nest.map_structure(
                  utils.merge_beam_dim, states
              )
              flat_logits, flat_states = symbols_to_logits_fn(
                  flat_ids, i, flat_states
              )
              states = nest.map_structure(
                  lambda t: utils.unmerge_beam_dim(
                      t, batch_size, beam_size
                  ),
                  flat_states,
              )
          else:
              flat_logits = symbols_to_logits_fn(flat_ids)
  
          logits = tf.reshape(flat_logits, [batch_size, beam_size, -1])
          candidate_log_probs = utils.log_prob_from_logits(logits)
          log_probs = candidate_log_probs + tf.expand_dims(
              alive_log_probs, axis=2
          )
  
          length_penalty = tf.pow(
              ((5.0 + tf.to_float(i + 1)) / 6.0), alpha
          )
  
          curr_scores = log_probs / length_penalty
          flat_curr_scores = tf.reshape(
              curr_scores, [-1, beam_size * vocab_size]
          )
  
          topk_scores, topk_ids = tf.nn.top_k(
              flat_curr_scores, k=beam_size * 2
          )
          topk_log_probs = topk_scores * length_penalty
  
          topk_beam_index = topk_ids // vocab_size
          topk_ids %= vocab_size  # Unflatten the ids
          batch_pos = utils.compute_batch_indices(
              batch_size, beam_size * 2
          )
          topk_coordinates = tf.stack(
              [batch_pos, topk_beam_index], axis=2
          )
  
          topk_seq = tf.gather_nd(alive_seq, topk_coordinates)
          if states:
              states = nest.map_structure(
                  lambda state: tf.gather_nd(state, topk_coordinates),
                  states,
              )
          topk_seq = tf.concat(
              [topk_seq, tf.expand_dims(topk_ids, axis=2)], axis=2
          )
  
          topk_finished = tf.equal(topk_ids, eos_id)
  
          return (
              topk_seq,
              topk_log_probs,
              topk_scores,
              topk_finished,
              states,
          )
  
      def inner_loop(
          i,
          alive_seq,
          alive_log_probs,
          finished_seq,
          finished_scores,
          finished_flags,
          states,
      ):
          """Inner beam search loop."""
          (
              topk_seq,
              topk_log_probs,
              topk_scores,
              topk_finished,
              states,
          ) = grow_topk(i, alive_seq, alive_log_probs, states)
          alive_seq, alive_log_probs, _, states = grow_alive(
              topk_seq,
              topk_scores,
              topk_log_probs,
              topk_finished,
              states,
          )
          (
              finished_seq,
              finished_scores,
              finished_flags,
              _,
          ) = grow_finished(
              finished_seq,
              finished_scores,
              finished_flags,
              topk_seq,
              topk_scores,
              topk_finished,
          )
  
          return (
              i + 1,
              alive_seq,
              alive_log_probs,
              finished_seq,
              finished_scores,
              finished_flags,
              states,
          )
  
      def _is_finished(
          i,
          unused_alive_seq,
          alive_log_probs,
          unused_finished_seq,
          finished_scores,
          unused_finished_in_finished,
          unused_states,
      ):
          """Checking termination condition."""
          max_length_penalty = tf.pow(
              ((5.0 + tf.to_float(decode_length)) / 6.0), alpha
          )
          lower_bound_alive_scores = (
              alive_log_probs[:, 0] / max_length_penalty
          )
  
          if not stop_early:
              lowest_score_of_finished_in_finished = tf.reduce_min(
                  finished_scores
              )
          else:
              lowest_score_of_finished_in_finished = tf.reduce_max(
                  finished_scores, axis=1
              )
  
          bound_is_met = tf.reduce_all(
              tf.greater(
                  lowest_score_of_finished_in_finished,
                  lower_bound_alive_scores,
              )
          )
  
          return tf.logical_and(
              tf.less(i, decode_length), tf.logical_not(bound_is_met)
          )
  
      inner_shape = tf.TensorShape([None, None, None])
  
      state_struc = nest.map_structure(
          utils.get_state_shape_invariants, states
      )
      (
          _,
          alive_seq,
          alive_log_probs,
          finished_seq,
          finished_scores,
          finished_flags,
          states,
      ) = tf.while_loop(
          _is_finished,
          inner_loop,
          [
              tf.constant(0),
              alive_seq,
              alive_log_probs,
              finished_seq,
              finished_scores,
              finished_flags,
              states,
          ],
          shape_invariants=[
              tf.TensorShape([]),
              inner_shape,
              alive_log_probs.get_shape(),
              inner_shape,
              finished_scores.get_shape(),
              finished_flags.get_shape(),
              state_struc,
          ],
          parallel_iterations=1,
          back_prop=False,
      )
  
      alive_seq.set_shape((None, beam_size, None))
      finished_seq.set_shape((None, beam_size, None))
      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,
      )
      return finished_seq, finished_scores, states
  
imports: import math
  from absl import logging
  import numpy
  import delta.compat
  from tensorflow.python.util import nest
  from delta.layers.base_layer import Layer
  from delta.layers import utils_tf
  import delta.layers
label_desc: Use image node_id 4 for calling the TransformerDecoder obj's underlying member method code with example usage: obj.beam_search(symbols_to_logits_fn, initial_ids, beam_size, decode_length, vocab_size, alpha, eos_id, states, stop_early, INF) and returns: finished_seq,  finished_scores,  states, utils, utils, topk_seq,  topk_log_probs,  topk_scores,  topk_finished,  states, unknown,  alive_seq,  alive_log_probs,  finished_seq,  finished_scores,  finished_flags,  states, tf