code: def test_step_with_scheduled_embedding_training_helper():
      sequence_length = [3, 4, 3, 1, 0]
      batch_size = 5
      max_time = 8
      input_depth = 7
      vocabulary_size = 10
  
      inputs = np.random.randn(
          batch_size, max_time, input_depth
      ).astype(np.float32)
      input_t = tf.constant(inputs)
      embeddings = np.random.randn(vocabulary_size, input_depth).astype(
          np.float32
      )
      half = tf.constant(0.5)
      cell = tf.keras.layers.LSTMCell(vocabulary_size)
      sampler = sampler_py.ScheduledEmbeddingTrainingSampler(
          sampling_probability=half, time_major=False
      )
      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,
          embedding=embeddings,
          initial_state=initial_state,
      )
      output_size = my_decoder.output_size
      output_dtype = my_decoder.output_dtype
      assert (
          basic_decoder.BasicDecoderOutput(
              vocabulary_size, 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 isinstance(step_outputs, basic_decoder.BasicDecoderOutput)
      assert (batch_size, vocabulary_size) == step_outputs[0].shape
      assert (batch_size,) == step_outputs[1].shape
      assert (batch_size, vocabulary_size) == first_state[0].shape
      assert (batch_size, vocabulary_size) == first_state[1].shape
      assert (batch_size, vocabulary_size) == step_state[0].shape
      assert (batch_size, vocabulary_size) == step_state[1].shape
      assert (batch_size, input_depth) == step_next_inputs.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.numpy()
      assert output_dtype.sample_id == sample_ids.dtype
      batch_where_not_sampling = np.where(sample_ids == -1)
      batch_where_sampling = np.where(sample_ids > -1)
  
      np.testing.assert_equal(
          step_next_inputs.numpy()[batch_where_sampling],
          embeddings[sample_ids[batch_where_sampling]],
      )
      np.testing.assert_equal(
          step_next_inputs.numpy()[batch_where_not_sampling],
          np.squeeze(inputs[batch_where_not_sampling, 1], axis=0),
      )
  
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 5 for calling a global function with example usage: test_step_with_scheduled_embedding_training_helper() without return types