code: def beam_search_step(time_, logits, beam_state, config):
      """Performs a single step of Beam Search Decoding.
  
      Args:
        time_: Beam search time step, should start at 0. At time 0 we assume
          that all beams are equal and consider only the first beam for
          continuations.
        logits: Logits at the current time step. A tensor of shape `[B, vocab_size]`
        beam_state: Current state of the beam search. An instance of `BeamState`
        config: An instance of `BeamSearchConfig`
  
      Returns:
        A new beam state.
      """
  
      # Calculate the current lengths of the predictions
      prediction_lengths = beam_state.lengths
      previously_finished = beam_state.finished
  
      # Calculate the total log probs for the new hypotheses
      # Final Shape: [beam_width, vocab_size]
      probs = tf.nn.log_softmax(logits)
      probs = mask_probs(probs, config.eos_token, previously_finished)
      total_probs = tf.expand_dims(beam_state.log_probs, 1) + probs
  
      # Calculate the continuation lengths
      # We add 1 to all continuations that are not EOS and were not
      # finished previously
      lengths_to_add = tf.one_hot(
          [config.eos_token] * config.beam_width,
          config.vocab_size,
          0,
          1,
      )
      add_mask = 1 - tf.to_int32(previously_finished)
      lengths_to_add = tf.expand_dims(add_mask, 1) * lengths_to_add
      new_prediction_lengths = (
          tf.expand_dims(prediction_lengths, 1) + lengths_to_add
      )
  
      # Calculate the scores for each beam
      scores = hyp_score(
          log_probs=total_probs,
          sequence_lengths=new_prediction_lengths,
          config=config,
      )
  
      scores_flat = tf.reshape(scores, [-1])
      # During the first time step we only consider the initial beam
      scores_flat = tf.cond(
          tf.convert_to_tensor(time_) > 0,
          lambda: scores_flat,
          lambda: scores[0],
      )
  
      # Pick the next beams according to the specified successors function
      next_beam_scores, word_indices = config.choose_successors_fn(
          scores_flat, config
      )
      next_beam_scores.set_shape([config.beam_width])
      word_indices.set_shape([config.beam_width])
  
      # Pick out the probs, beam_ids, and states according to the chosen predictions
      total_probs_flat = tf.reshape(
          total_probs, [-1], name="total_probs_flat"
      )
      next_beam_probs = tf.gather(total_probs_flat, word_indices)
      next_beam_probs.set_shape([config.beam_width])
      next_word_ids = tf.mod(word_indices, config.vocab_size)
      next_beam_ids = tf.div(word_indices, config.vocab_size)
  
      # Append new ids to current predictions
      next_finished = tf.logical_or(
          tf.gather(beam_state.finished, next_beam_ids),
          tf.equal(next_word_ids, config.eos_token),
      )
  
      # Calculate the length of the next predictions.
      # 1. Finished beams remain unchanged
      # 2. Beams that are now finished (EOS predicted) remain unchanged
      # 3. Beams that are not yet finished have their length increased by 1
      lengths_to_add = tf.to_int32(
          tf.not_equal(next_word_ids, config.eos_token)
      )
      lengths_to_add = (1 - tf.to_int32(next_finished)) * lengths_to_add
      next_prediction_len = tf.gather(beam_state.lengths, next_beam_ids)
      next_prediction_len += lengths_to_add
  
      next_state = BeamSearchState(
          log_probs=next_beam_probs,
          lengths=next_prediction_len,
          finished=next_finished,
      )
  
      output = BeamSearchStepOutput(
          scores=next_beam_scores,
          predicted_ids=next_word_ids,
          beam_parent_ids=next_beam_ids,
      )
  
      return output, next_state
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from collections import namedtuple
  import numpy
  import tensorflow
  from tensorflow.python.util import nest
label_desc: Use image node_id 9 for calling a global function with example usage: beam_search_step(time_, logits, beam_state, config) and returns: output,  next_state