code: def test_eos_masking():
      probs = tf.constant(
          [
              [
                  [-0.2, -0.2, -0.2, -0.2, -0.2],
                  [-0.3, -0.3, -0.3, 3, 0],
                  [5, 6, 0, 0, 0],
              ],
              [
                  [-0.2, -0.2, -0.2, -0.2, 0],
                  [-0.3, -0.3, -0.1, 3, 0],
                  [5, 6, 3, 0, 0],
              ],
          ]
      )
  
      eos_token = 0
      previously_finished = np.array([[0, 1, 0], [0, 1, 1]], dtype=bool)
      masked = beam_search_decoder._mask_probs(
          probs, eos_token, previously_finished
      )
      masked = masked.numpy()
  
      np.testing.assert_equal(probs[0][0], masked[0][0])
      np.testing.assert_equal(probs[0][2], masked[0][2])
      np.testing.assert_equal(probs[1][0], masked[1][0])
  
      np.testing.assert_equal(masked[0][1][0], 0)
      np.testing.assert_equal(masked[1][1][0], 0)
      np.testing.assert_equal(masked[1][2][0], 0)
  
      for i in range(1, 5):
          np.testing.assert_allclose(
              masked[0][1][i], np.finfo("float32").min
          )
          np.testing.assert_allclose(
              masked[1][1][i], np.finfo("float32").min
          )
          np.testing.assert_allclose(
              masked[1][2][i], np.finfo("float32").min
          )
  
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 11 for calling a global function with example usage: test_eos_masking() without return types