code: def _beam_search_step(
      time,
      logits,
      next_cell_state,
      beam_state,
      batch_size,
      beam_width,
      end_token,
      length_penalty_weight,
      coverage_penalty_weight,
      output_all_scores,
  ):
      """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
          `[batch_size, beam_width, vocab_size]`
        next_cell_state: The next state from the cell, e.g. an instance of
          AttentionWrapperState if the cell is attentional.
        beam_state: Current state of the beam search.
          An instance of `BeamSearchDecoderState`.
        batch_size: The batch size for this input.
        beam_width: Python int.  The size of the beams.
        end_token: The int32 end token.
        length_penalty_weight: Float weight to penalize length. Disabled with
          0.0.
        coverage_penalty_weight: Float weight to penalize the coverage of source
          sentence. Disabled with 0.0.
        output_all_scores: Bool output scores for every token if True, else only
          output the top scores.
  
      Returns:
        A new beam state.
      """
      static_batch_size = tf.get_static_value(batch_size)
  
      # Calculate the current lengths of the predictions
      prediction_lengths = beam_state.lengths
      previously_finished = beam_state.finished
      not_finished = tf.logical_not(previously_finished)
  
      # Calculate the total log probs for the new hypotheses
      # Final Shape: [batch_size, beam_width, vocab_size]
      step_log_probs = tf.nn.log_softmax(logits)
      step_log_probs = _mask_probs(
          step_log_probs, end_token, previously_finished
      )
      total_probs = (
          tf.expand_dims(beam_state.log_probs, 2) + step_log_probs
      )
  
      # Calculate the continuation lengths by adding to all continuing beams.
      vocab_size = logits.shape[-1] or tf.shape(logits)[-1]
      lengths_to_add = tf.one_hot(
          indices=tf.fill([batch_size, beam_width], end_token),
          depth=vocab_size,
          on_value=np.int64(0),
          off_value=np.int64(1),
          dtype=tf.int64,
      )
      add_mask = tf.cast(not_finished, tf.int64)
      lengths_to_add *= tf.expand_dims(add_mask, 2)
      new_prediction_lengths = lengths_to_add + tf.expand_dims(
          prediction_lengths, 2
      )
  
      # Calculate the accumulated attention probabilities if coverage penalty is
      # enabled.
      accumulated_attention_probs = None
      attention_probs = get_attention_probs(
          next_cell_state, coverage_penalty_weight
      )
      if attention_probs is not None:
          attention_probs *= tf.expand_dims(
              tf.cast(not_finished, tf.float32), 2
          )
          accumulated_attention_probs = (
              beam_state.accumulated_attention_probs + attention_probs
          )
  
      # Calculate the scores for each beam
      scores = _get_scores(
          log_probs=total_probs,
          sequence_lengths=new_prediction_lengths,
          length_penalty_weight=length_penalty_weight,
          coverage_penalty_weight=coverage_penalty_weight,
          finished=previously_finished,
          accumulated_attention_probs=accumulated_attention_probs,
      )
  
      time = tf.convert_to_tensor(time, name="time")
      # During the first time step we only consider the initial beam
      scores_flat = tf.reshape(scores, [batch_size, -1])
  
      # Pick the next beams according to the specified successors function
      next_beam_size = tf.convert_to_tensor(
          beam_width, dtype=tf.int32, name="beam_width"
      )
      next_beam_scores, word_indices = tf.math.top_k(
          scores_flat, k=next_beam_size
      )
  
      next_beam_scores.set_shape([static_batch_size, beam_width])
      word_indices.set_shape([static_batch_size, beam_width])
  
      # Pick out the probs, beam_ids, and states according to the chosen
      # predictions
      next_beam_probs = _tensor_gather_helper(
          gather_indices=word_indices,
          gather_from=total_probs,
          batch_size=batch_size,
          range_size=beam_width * vocab_size,
          gather_shape=[-1],
          name="next_beam_probs",
      )
      # Note: just doing the following
      #   tf.to_int32(word_indices % vocab_size,
      #       name="next_beam_word_ids")
      # would be a lot cleaner but for reasons unclear, that hides the results of
      # the op which prevents capturing it with tfdbg debug ops.
      raw_next_word_ids = tf.math.floormod(
          word_indices, vocab_size, name="next_beam_word_ids"
      )
      next_word_ids = tf.cast(raw_next_word_ids, tf.int32)
      next_beam_ids = tf.cast(
          word_indices / vocab_size,
          tf.int32,
          name="next_beam_parent_ids",
      )
  
      # Append new ids to current predictions
      previously_finished = _tensor_gather_helper(
          gather_indices=next_beam_ids,
          gather_from=previously_finished,
          batch_size=batch_size,
          range_size=beam_width,
          gather_shape=[-1],
      )
      next_finished = tf.logical_or(
          previously_finished,
          tf.equal(next_word_ids, end_token),
          name="next_beam_finished",
      )
  
      # Calculate the length of the next predictions.
      # 1. Finished beams remain unchanged.
      # 2. Beams that are now finished (EOS predicted) have their length
      #    increased by 1.
      # 3. Beams that are not yet finished have their length increased by 1.
      lengths_to_add = tf.cast(
          tf.logical_not(previously_finished), tf.int64
      )
      next_prediction_len = _tensor_gather_helper(
          gather_indices=next_beam_ids,
          gather_from=beam_state.lengths,
          batch_size=batch_size,
          range_size=beam_width,
          gather_shape=[-1],
      )
      next_prediction_len += lengths_to_add
      next_accumulated_attention_probs = ()
      if accumulated_attention_probs is not None:
          next_accumulated_attention_probs = _tensor_gather_helper(
              gather_indices=next_beam_ids,
              gather_from=accumulated_attention_probs,
              batch_size=batch_size,
              range_size=beam_width,
              gather_shape=[batch_size * beam_width, -1],
              name="next_accumulated_attention_probs",
          )
  
      # Pick out the cell_states according to the next_beam_ids. We use a
      # different gather_shape here because the cell_state tensors, i.e.
      # the tensors that would be gathered from, all have dimension
      # greater than two and we need to preserve those dimensions.
      next_cell_state = tf.nest.map_structure(
          lambda gather_from: _maybe_tensor_gather_helper(
              gather_indices=next_beam_ids,
              gather_from=gather_from,
              batch_size=batch_size,
              range_size=beam_width,
              gather_shape=[batch_size * beam_width, -1],
          ),
          next_cell_state,
      )
  
      next_state = BeamSearchDecoderState(
          cell_state=next_cell_state,
          log_probs=next_beam_probs,
          lengths=next_prediction_len,
          finished=next_finished,
          accumulated_attention_probs=next_accumulated_attention_probs,
      )
  
      output = BeamSearchDecoderOutput(
          scores=scores if output_all_scores else next_beam_scores,
          predicted_ids=next_word_ids,
          parent_ids=next_beam_ids,
      )
  
      return output, next_state
  
imports: import collections
  import numpy
  import tensorflow
  from tensorflow_addons import options
  from tensorflow_addons.seq2seq import attention_wrapper
  from tensorflow_addons.seq2seq import decoder
  from tensorflow_addons.utils import keras_utils
  from tensorflow_addons.utils.resource_loader import LazySO
  from tensorflow_addons.utils.types import FloatTensorLike, TensorLike, Number
  from typeguard import typechecked
  from typing import Callable, Optional
label_desc: Use image node_id 10 for calling a global function with example usage: _beam_search_step(time, logits, next_cell_state, beam_state, batch_size, beam_width, end_token, length_penalty_weight, coverage_penalty_weight, output_all_scores) and returns: output,  next_state