code: def test_step_with_eos():
      batch_size = 2
      beam_width = 3
      vocab_size = 5
      end_token = 0
      length_penalty_weight = 0.6
      coverage_penalty_weight = 0.0
      output_all_scores = False
  
      dummy_cell_state = tf.zeros([batch_size, beam_width])
      beam_state = beam_search_decoder.BeamSearchDecoderState(
          cell_state=dummy_cell_state,
          log_probs=tf.nn.log_softmax(
              tf.ones([batch_size, beam_width])
          ),
          lengths=tf.convert_to_tensor(
              [[2, 1, 2], [2, 2, 1]], dtype=tf.int64
          ),
          finished=tf.convert_to_tensor(
              [[False, True, False], [False, False, True]],
              dtype=tf.bool,
          ),
          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] = 5.7  # why does this not work when it's 2.7?
      logits_[1, 2, 2] = 1.0
      logits_[1, 2, 3] = 0.2
      logits = tf.convert_to_tensor(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_, state_, log_probs_ = [
          outputs,
          next_beam_state,
          beam_state,
          log_probs,
      ]
  
      np.testing.assert_equal(
          outputs_.parent_ids.numpy(),
          np.asanyarray([[1, 0, 0], [1, 2, 0]]),
      )
      np.testing.assert_equal(
          outputs_.predicted_ids.numpy(),
          np.asanyarray([[0, 3, 2], [2, 0, 1]]),
      )
      np.testing.assert_equal(
          next_state_.lengths.numpy(),
          np.asanyarray([[1, 3, 3], [3, 1, 3]]),
      )
      np.testing.assert_equal(
          next_state_.finished.numpy(),
          np.asanyarray([[True, False, False], [False, True, False]]),
      )
  
      expected_log_probs = []
      expected_log_probs.append(state_.log_probs[0].numpy())
      expected_log_probs.append(state_.log_probs[1].numpy())
      expected_log_probs[0][1] += log_probs_[0, 0, 3]
      expected_log_probs[0][2] += log_probs_[0, 0, 2]
      expected_log_probs[1][0] += log_probs_[1, 1, 2]
      expected_log_probs[1][2] += log_probs_[1, 0, 1]
      np.testing.assert_equal(
          next_state_.log_probs.numpy(),
          np.asanyarray(expected_log_probs),
      )
  
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 14 for calling a global function with example usage: test_step_with_eos() without return types