code: def testNotGreedyBatchTwoBeamTwoWithAlpha(self):
      batch_size = 2
      beam_size = 2
      vocab_size = 3
      decode_length = 3
  
      initial_ids = tf.constant([0] * batch_size)  # GO
      # Probabilities for position * batch * beam * vocab
      # Probabilities have been set such that with alpha = 3.5, the less probable
      # but longer sequence will have a better score than the shorter sequence
      # with higher log prob in batch 1, and the order will be reverse in batch
      # 2. That is, the shorter sequence will still have a higher score in spite
      # of the length penalty
      probabilities = tf.constant(
          [
              [
                  [[0.1, 0.1, 0.8], [0.1, 0.1, 0.8]],
                  [[0.1, 0.1, 0.8], [0.1, 0.1, 0.8]],
              ],
              [
                  [[0.4, 0.5, 0.1], [0.2, 0.4, 0.4]],
                  [[0.3, 0.6, 0.1], [0.2, 0.4, 0.4]],
              ],
              [
                  [[0.05, 0.9, 0.05], [0.4, 0.4, 0.2]],
                  [[0.05, 0.9, 0.05], [0.4, 0.4, 0.2]],
              ],
          ]
      )
  
      def symbols_to_logits(ids):
          pos = tf.shape(ids)[1]
          logits = tf.to_float(tf.log(probabilities[pos - 1, :]))
          return logits
  
      final_ids, final_scores, _ = beam_search.beam_search(
          symbols_to_logits,
          initial_ids,
          beam_size,
          decode_length,
          vocab_size,
          3.5,
          eos_id=1,
      )
  
      with self.test_session():
          ids = final_ids.eval()
          scores = final_scores.eval()
      self.assertAllEqual(
          [[[0, 2, 0, 1], [0, 2, 1, 0]], [[0, 2, 1, 0], [0, 2, 0, 1]]],
          ids,
      )
      self.assertAllClose(
          [
              [
                  np.log(0.8 * 0.4 * 0.9) / (8.0 / 6.0) ** 3.5,
                  np.log(0.8 * 0.5) / (7.0 / 6.0) ** 3.5,
              ],
              [
                  np.log(0.8 * 0.6) / (7.0 / 6.0) ** 3.5,
                  np.log(0.8 * 0.3 * 0.9) / (8.0 / 6.0) ** 3.5,
              ],
          ],
          scores,
      )
  
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 7 for calling the BeamSearchTest obj's underlying member method code with example usage: obj.testNotGreedyBatchTwoBeamTwoWithAlpha() and returns: logits