code: def testStates(self):
      batch_size = 1
      beam_size = 1
      vocab_size = 2
      decode_length = 3
  
      initial_ids = tf.constant([0] * batch_size)  # GO
      probabilities = tf.constant(
          [[[0.7, 0.3]], [[0.4, 0.6]], [[0.5, 0.5]]]
      )
  
      expected_states = tf.constant([[[0.0]], [[1.0]]])
  
      def symbols_to_logits(ids, _, states):
          pos = tf.shape(ids)[1] - 1
          # We have to assert the values of state inline here since we can't fetch
          # them out of the loop!
          with tf.control_dependencies(
              [tf.assert_equal(states["state"], expected_states[pos])]
          ):
              logits = tf.to_float(tf.log(probabilities[pos, :]))
  
          states["state"] += 1
          return logits, states
  
      states = {
          "state": tf.zeros((batch_size, 1)),
      }
      states["state"] = tf.placeholder_with_default(
          states["state"], shape=(None, 1)
      )
  
      final_ids, _, _ = beam_search.beam_search(
          symbols_to_logits,
          initial_ids,
          beam_size,
          decode_length,
          vocab_size,
          0.0,
          eos_id=1,
          states=states,
      )
  
      with self.test_session() as sess:
          # Catch and fail so that the testing framework doesn't think it's an error
          try:
              sess.run(final_ids)
          except tf.errors.InvalidArgumentError as e:
              raise AssertionError(e.message)
  
file_path: tensor2tensor/tensor2tensor/utils/beam_search_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import numpy
  from tensor2tensor.utils import beam_search
  import tensorflow.compat.v1
label_desc: Use image node_id 9 for calling the BeamSearchTest obj's underlying member method code with example usage: obj.testStates() and returns: logits,  states