code: def test_step_with_scheduled_output_training_helper(
      sampling_probability, use_next_inputs_fn, use_auxiliary_inputs
  ):
      sequence_length = [3, 4, 3, 1, 0]
      batch_size = 5
      max_time = 8
      input_depth = 7
      cell_depth = input_depth
      if use_auxiliary_inputs:
          auxiliary_input_depth = 4
          auxiliary_inputs = np.random.randn(
              batch_size, max_time, auxiliary_input_depth
          ).astype(np.float32)
      else:
          auxiliary_inputs = None
  
      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)
      sampling_probability = tf.constant(sampling_probability)
  
      if use_next_inputs_fn:
  
          def next_inputs_fn(outputs):
              # Use deterministic function for test.
              samples = tf.argmax(outputs, axis=1)
              return tf.one_hot(samples, cell_depth, dtype=tf.float32)
  
      else:
          next_inputs_fn = None
  
      sampler = sampler_py.ScheduledOutputTrainingSampler(
          sampling_probability=sampling_probability,
          time_major=False,
          next_inputs_fn=next_inputs_fn,
      )
      initial_state = cell.get_initial_state(
          batch_size=batch_size, dtype=tf.float32
      )
  
      my_decoder = basic_decoder.BasicDecoder(
          cell=cell, sampler=sampler
      )
  
      (
          first_finished,
          first_inputs,
          first_state,
      ) = my_decoder.initialize(
          input_t,
          sequence_length=sequence_length,
          initial_state=initial_state,
          auxiliary_inputs=auxiliary_inputs,
      )
      output_size = my_decoder.output_size
      output_dtype = my_decoder.output_dtype
      assert (
          basic_decoder.BasicDecoderOutput(
              cell_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 use_next_inputs_fn:
          output_after_next_inputs_fn = next_inputs_fn(
              step_outputs.rnn_output
          )
  
      assert len(first_state) == 2
      assert len(step_state) == 2
      assert isinstance(step_outputs, basic_decoder.BasicDecoderOutput)
      assert (batch_size, cell_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
  
      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,
      )
  
      sample_ids = step_outputs.sample_id
      assert output_dtype.sample_id == sample_ids.dtype
      batch_where_not_sampling = np.where(np.logical_not(sample_ids))
      batch_where_sampling = np.where(sample_ids)
  
      auxiliary_inputs_to_concat = (
          auxiliary_inputs[:, 1]
          if use_auxiliary_inputs
          else np.array([]).reshape(batch_size, 0).astype(np.float32)
      )
  
      expected_next_sampling_inputs = np.concatenate(
          (
              output_after_next_inputs_fn.numpy()[batch_where_sampling]
              if use_next_inputs_fn
              else step_outputs.rnn_output.numpy()[
                  batch_where_sampling
              ],
              auxiliary_inputs_to_concat[batch_where_sampling],
          ),
          axis=-1,
      )
  
      np.testing.assert_equal(
          step_next_inputs.numpy()[batch_where_sampling],
          expected_next_sampling_inputs,
      )
  
      np.testing.assert_equal(
          step_next_inputs.numpy()[batch_where_not_sampling],
          np.concatenate(
              (
                  np.squeeze(
                      inputs[batch_where_not_sampling, 1], axis=0
                  ),
                  auxiliary_inputs_to_concat[batch_where_not_sampling],
              ),
              axis=-1,
          ),
      )
  
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 6 for calling a global function with example usage: test_step_with_scheduled_output_training_helper(sampling_probability, use_next_inputs_fn, use_auxiliary_inputs) and returns: tf