code: def test_step_with_training_helper_output_layer(
      cell_class, use_output_layer
  ):
      sequence_length = [3, 4, 3, 1, 0]
      batch_size = 5
      max_time = 8
      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 = cell_class(cell_depth)
      sampler = sampler_py.TrainingSampler(time_major=False)
      if use_output_layer:
          output_layer = tf.keras.layers.Dense(
              output_layer_depth, use_bias=False
          )
          expected_output_depth = output_layer_depth
      else:
          output_layer = None
          expected_output_depth = cell_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
      )
  
      (
          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)
  
      if isinstance(cell, tf.keras.layers.LSTMCell):
          assert len(first_state) == 2
          assert len(step_state) == 2
          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
      elif isinstance(cell, tf.keras.layers.GRUCell):
          assert tf.is_tensor(first_state)
          assert tf.is_tensor(step_state)
          assert (batch_size, cell_depth) == first_state.shape
          assert (batch_size, cell_depth) == step_state.shape
      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
  
      if use_output_layer:
          # The output layer was accessed
          assert len(output_layer.variables) == 1
  
      np.testing.assert_equal(
          np.asanyarray([False, False, False, False, True]),
          first_finished,
      )
      np.testing.assert_equal(
          np.asanyarray([False, False, False, True, True]),
          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 1 for calling a global function with example usage: test_step_with_training_helper_output_layer(cell_class, use_output_layer) without return types