code: def test_large_beam_step():
      batch_size = 2
      beam_width = 8
      vocab_size = 5
      end_token = 0
      length_penalty_weight = 0.6
      coverage_penalty_weight = 0.0
      output_all_scores = False
  
      def get_probs():
          """this simulates the initialize method in BeamSearchDecoder."""
          log_prob_mask = tf.one_hot(
              tf.zeros([batch_size], dtype=tf.int32),
              depth=beam_width,
              on_value=True,
              off_value=False,
              dtype=tf.bool,
          )
  
          log_prob_zeros = tf.zeros(
              [batch_size, beam_width], dtype=tf.float32
          )
          log_prob_neg_inf = (
              tf.ones([batch_size, beam_width], dtype=tf.float32)
              * -np.Inf
          )
  
          log_probs = tf.where(
              log_prob_mask, log_prob_zeros, log_prob_neg_inf
          )
          return log_probs
  
      log_probs = get_probs()
      dummy_cell_state = tf.zeros([batch_size, beam_width])
  
      _finished = tf.one_hot(
          tf.zeros([batch_size], dtype=tf.int32),
          depth=beam_width,
          on_value=False,
          off_value=True,
          dtype=tf.bool,
      )
      _lengths = np.zeros([batch_size, beam_width], dtype=np.int64)
      _lengths[:, 0] = 2
      _lengths = tf.constant(_lengths, dtype=tf.int64)
  
      beam_state = beam_search_decoder.BeamSearchDecoderState(
          cell_state=dummy_cell_state,
          log_probs=log_probs,
          lengths=_lengths,
          finished=_finished,
          accumulated_attention_probs=(),
      )
  
      logits_ = np.full([batch_size, beam_width, vocab_size], 0.0001)
      logits_[0, 0, 2] = 1.9
      logits_[0, 0, 3] = 2.1
      logits_[0, 1, 3] = 3.1
      logits_[0, 1, 4] = 0.9
      logits_[1, 0, 1] = 0.5
      logits_[1, 1, 2] = 2.7
      logits_[1, 2, 2] = 10.0
      logits_[1, 2, 3] = 0.2
      logits = tf.constant(logits_, dtype=tf.float32)
      log_probs = tf.nn.log_softmax(logits)
  
      outputs, next_beam_state = beam_search_decoder._beam_search_step(
          time=2,
          logits=logits,
          next_cell_state=dummy_cell_state,
          beam_state=beam_state,
          batch_size=tf.convert_to_tensor(batch_size),
          beam_width=beam_width,
          end_token=end_token,
          length_penalty_weight=length_penalty_weight,
          coverage_penalty_weight=coverage_penalty_weight,
          output_all_scores=output_all_scores,
      )
  
      outputs_, next_state_ = [outputs, next_beam_state]
  
      assert outputs_.predicted_ids[0, 0] == 3
      assert outputs_.predicted_ids[0, 1] == 2
      assert outputs_.predicted_ids[1, 0] == 1
      neg_inf = -np.Inf
      np.testing.assert_equal(
          next_state_.log_probs[:, -3:].numpy(),
          np.asanyarray(
              [[neg_inf, neg_inf, neg_inf], [neg_inf, neg_inf, neg_inf]]
          ),
      )
      np.testing.assert_equal(
          np.asanyarray(next_state_.log_probs[:, :-3] > neg_inf), True
      )
      np.testing.assert_equal(
          np.asanyarray(next_state_.lengths[:, :-3] > 0), True
      )
      np.testing.assert_equal(
          next_state_.lengths[:, -3:].numpy(),
          np.asanyarray([[0, 0, 0], [0, 0, 0]]),
      )
  
imports: import numpy
  import pytest
  import tensorflow
  from tensorflow_addons.seq2seq import attention_wrapper
  from tensorflow_addons.seq2seq import beam_search_decoder, gather_tree
label_desc: Use image node_id 15 for calling a global function with example usage: test_large_beam_step() and returns: log_probs