code: def testShapes(self):
      batch_size = 2
      beam_size = 3
      vocab_size = 4
      decode_length = 10
  
      initial_ids = tf.constant([0, 0])  # GO
  
      def symbols_to_logits(_):
          # Just return random logits
          return tf.random_uniform((batch_size * beam_size, vocab_size))
  
      final_ids, final_probs, _ = beam_search.beam_search(
          symbols_to_logits,
          initial_ids,
          beam_size,
          decode_length,
          vocab_size,
          0.0,
      )
  
      self.assertEqual(
          final_ids.get_shape().as_list(), [None, beam_size, None]
      )
  
      self.assertEqual(
          final_probs.get_shape().as_list(), [batch_size, beam_size]
      )
  
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 1 for calling the BeamSearchTest obj's underlying member method code with example usage: obj.testShapes() and returns: tf