code: def test_step_with_training_helper_masked_input(use_mask):
      batch_size = 5
      max_time = 8
      sequence_length = (
          [max_time] * batch_size
          if use_mask is None
          else [3, 4, 3, 1, 0]
      )
      sequence_length = np.array(sequence_length, dtype=np.int32)
      mask = [
          [True] * sl + [False] * (max_time - sl)
          for sl in sequence_length
      ]
      input_depth = 7
      cell_depth = 10
      output_layer_depth = 3
  
      inputs = np.random.randn(
          batch_size, max_time, input_depth
      ).astype(np.float32)
      input_t = tf.constant(inputs)
      cell = tf.keras.layers.LSTMCell(cell_depth)
      sampler = sampler_py.TrainingSampler(time_major=False)
      output_layer = tf.keras.layers.Dense(
          output_layer_depth, use_bias=False
      )
      expected_output_depth = output_layer_depth
      initial_state = cell.get_initial_state(
          batch_size=batch_size, dtype=tf.float32
      )
      my_decoder = basic_decoder.BasicDecoder(
          cell=cell, sampler=sampler, output_layer=output_layer
      )
  
      if use_mask is None:
          (
              first_finished,
              first_inputs,
              first_state,
          ) = my_decoder.initialize(
              input_t, initial_state=initial_state
          )
      elif use_mask:
          (
              first_finished,
              first_inputs,
              first_state,
          ) = my_decoder.initialize(
              input_t, initial_state=initial_state, mask=mask
          )
      else:
          (
              first_finished,
              first_inputs,
              first_state,
          ) = my_decoder.initialize(
              input_t,
              initial_state=initial_state,
              sequence_length=sequence_length,
          )
  
      output_size = my_decoder.output_size
      output_dtype = my_decoder.output_dtype
      assert (
          basic_decoder.BasicDecoderOutput(
              expected_output_depth, tf.TensorShape([])
          )
          == output_size
      )
  
      assert (
          basic_decoder.BasicDecoderOutput(tf.float32, tf.int32)
          == output_dtype
      )
  
      (
          step_outputs,
          step_state,
          step_next_inputs,
          step_finished,
      ) = my_decoder.step(tf.constant(0), first_inputs, first_state)
  
      assert len(first_state) == 2
      assert len(step_state) == 2
      assert type(step_outputs) is basic_decoder.BasicDecoderOutput
      assert (batch_size, expected_output_depth) == step_outputs[
          0
      ].shape
      assert (batch_size,) == step_outputs[1].shape
      assert (batch_size, cell_depth) == first_state[0].shape
      assert (batch_size, cell_depth) == first_state[1].shape
      assert (batch_size, cell_depth) == step_state[0].shape
      assert (batch_size, cell_depth) == step_state[1].shape
  
      assert len(output_layer.variables) == 1
  
      np.testing.assert_equal(sequence_length == 0, first_finished)
      np.testing.assert_equal(
          (np.maximum(sequence_length - 1, 0) == 0), step_finished
      )
      assert output_dtype.sample_id == step_outputs.sample_id.dtype
      np.testing.assert_equal(
          np.argmax(step_outputs.rnn_output, -1), step_outputs.sample_id
      )
  
file_path: addons/tensorflow_addons/seq2seq/tests/basic_decoder_test.py
imports: import numpy
  import pytest
  import tensorflow
  from tensorflow_addons.seq2seq import attention_wrapper
  from tensorflow_addons.seq2seq import basic_decoder
  from tensorflow_addons.seq2seq import sampler
label_desc: Use image node_id 2 for calling a global function with example usage: test_step_with_training_helper_masked_input(use_mask) without return types