code: def test_step_with_inference_helper_categorical():
      batch_size = 5
      vocabulary_size = 7
      cell_depth = vocabulary_size
      start_token = 0
      end_token = 6
  
      start_inputs = tf.one_hot(
          np.ones(batch_size, dtype=np.int32) * start_token,
          vocabulary_size,
      )
  
      # The sample function samples categorically from the logits.
      def sample_fn(x):
          return sampler_py.categorical_sample(logits=x)
  
      # The next inputs are a one-hot encoding of the sampled labels.
      def next_inputs_fn(x):
          return tf.one_hot(x, vocabulary_size, dtype=tf.float32)
  
      def end_fn(sample_ids):
          return tf.equal(sample_ids, end_token)
  
      cell = tf.keras.layers.LSTMCell(vocabulary_size)
      sampler = sampler_py.InferenceSampler(
          sample_fn,
          sample_shape=(),
          sample_dtype=tf.int32,
          end_fn=end_fn,
          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(
          start_inputs, initial_state=initial_state
      )
  
      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)
  
      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
  
      sample_ids = step_outputs.sample_id.numpy()
      assert output_dtype.sample_id == sample_ids.dtype
      expected_step_finished = sample_ids == end_token
      expected_step_next_inputs = np.zeros(
          (batch_size, vocabulary_size)
      )
      expected_step_next_inputs[np.arange(batch_size), sample_ids] = 1.0
      np.testing.assert_equal(expected_step_finished, step_finished)
      np.testing.assert_equal(
          expected_step_next_inputs, step_next_inputs
      )
  
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 7 for calling a global function with example usage: test_step_with_inference_helper_categorical() and returns: sampler_py, tf, tf