code: def testGreedyBatchOne(self):
      batch_size = 1
      beam_size = 1
      vocab_size = 2
      decode_length = 3
  
      initial_ids = tf.constant([0] * batch_size)  # GO
  
      # Test that beam search finds the most probable sequence.
      # These probabilities represent the following search
      #
      #               G0 (0)
      #                  / \
      #                /     \
      #              /         \
      #            /             \
      #         0(0.7)          1(0.3)
      #           / \
      #          /   \
      #         /     \
      #     0(0.4) 1(0.6)
      #        /\
      #       /  \
      #      /    \
      #    0(0.5) 1(0.5)
      # and the following decoding probabilities
      # 0000 - 0.7 * 0.4  * 0.1
      # 0001 - 0.7 * 0.4  * 0.9
      # 001 - 0.7 * 0.6 (Best)
      # 01 = 0.3
      #
      # 001 is the most likely sequence under these probabilities.
      probabilities = tf.constant(
          [[[0.7, 0.3]], [[0.4, 0.6]], [[0.5, 0.5]]]
      )
  
      def symbols_to_logits(ids):
          pos = tf.shape(ids)[1]
          logits = tf.to_float(tf.log(probabilities[pos - 1, :]))
          return logits
  
      final_ids, final_probs, _ = beam_search.beam_search(
          symbols_to_logits,
          initial_ids,
          beam_size,
          decode_length,
          vocab_size,
          0.0,
          eos_id=1,
      )
  
      with self.test_session():
          ids = final_ids.eval()
          probs = final_probs.eval()
      self.assertAllEqual([[[0, 0, 1]]], ids)
      self.assertAllClose([[0.7 * 0.6]], np.exp(probs))
  
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 3 for calling the BeamSearchTest obj's underlying member method code with example usage: obj.testGreedyBatchOne() and returns: logits