code: def test_beam_search_decoder(
      cell_class,
      time_major,
      has_attention,
      with_alignment_history,
      output_all_scores,
  ):
      encoder_sequence_length = np.array([3, 2, 3, 1, 1])
      batch_size = 5
      decoder_max_time = 4
      input_depth = 7
      cell_depth = 9
      attention_depth = 6
      vocab_size = 20
      end_token = vocab_size - 1
      start_token = 0
      embedding_dim = 50
      maximum_iterations = 3
      output_layer = tf.keras.layers.Dense(
          vocab_size, use_bias=True, activation=None
      )
      beam_width = 3
      embedding = tf.random.normal([vocab_size, embedding_dim])
      cell = cell_class(cell_depth)
  
      if has_attention:
          attention_mechanism = attention_wrapper.BahdanauAttention(
              units=attention_depth,
          )
          cell = attention_wrapper.AttentionWrapper(
              cell=cell,
              attention_mechanism=attention_mechanism,
              attention_layer_size=attention_depth,
              alignment_history=with_alignment_history,
          )
          coverage_penalty_weight = 0.2
      else:
          coverage_penalty_weight = 0.0
  
      bsd = beam_search_decoder.BeamSearchDecoder(
          cell=cell,
          beam_width=beam_width,
          output_layer=output_layer,
          length_penalty_weight=0.0,
          coverage_penalty_weight=coverage_penalty_weight,
          output_time_major=time_major,
          maximum_iterations=maximum_iterations,
          output_all_scores=output_all_scores,
      )
  
      @tf.function(
          input_signature=(
              tf.TensorSpec(
                  [None, None, input_depth], dtype=tf.float32
              ),
              tf.TensorSpec([None], dtype=tf.int32),
          )
      )
      def _beam_decode_from(memory, memory_sequence_length):
          batch_size_tensor = tf.shape(memory)[0]
  
          if has_attention:
              tiled_memory = beam_search_decoder.tile_batch(
                  memory, multiplier=beam_width
              )
              tiled_memory_sequence_length = (
                  beam_search_decoder.tile_batch(
                      memory_sequence_length, multiplier=beam_width
                  )
              )
              attention_mechanism.setup_memory(
                  tiled_memory,
                  memory_sequence_length=tiled_memory_sequence_length,
              )
  
          cell_state = cell.get_initial_state(
              batch_size=batch_size_tensor * beam_width,
              dtype=tf.float32,
          )
  
          return bsd(
              embedding,
              start_tokens=tf.fill([batch_size_tensor], start_token),
              end_token=end_token,
              initial_state=cell_state,
          )
  
      memory = tf.random.normal(
          [batch_size, decoder_max_time, input_depth]
      )
      memory_sequence_length = tf.constant(
          encoder_sequence_length, dtype=tf.int32
      )
      (
          final_outputs,
          final_state,
          final_sequence_lengths,
      ) = _beam_decode_from(memory, memory_sequence_length)
  
      def _t(shape):
          if time_major:
              return (shape[1], shape[0]) + shape[2:]
          return shape
  
      assert isinstance(
          final_outputs,
          beam_search_decoder.FinalBeamSearchDecoderOutput,
      )
      assert isinstance(
          final_state, beam_search_decoder.BeamSearchDecoderState
      )
  
      beam_search_decoder_output = (
          final_outputs.beam_search_decoder_output
      )
      max_sequence_length = np.max(final_sequence_lengths.numpy())
      assert _t((batch_size, max_sequence_length, beam_width)) == tuple(
          final_outputs.predicted_ids.shape.as_list()
      )
  
      if output_all_scores:
          assert _t(
              (batch_size, max_sequence_length, beam_width, vocab_size)
          ) == tuple(beam_search_decoder_output.scores.shape.as_list())
  
          # Check that the vocab size corresponds to the dimensions of the output.
          assert (beam_width, vocab_size) == tuple(
              bsd.output_size.scores.as_list()
          )
      else:
          assert _t(
              (batch_size, max_sequence_length, beam_width)
          ) == tuple(beam_search_decoder_output.scores.shape.as_list())
  
          # Check only the beam width corresponds to the dimensions of the output.
          assert (beam_width,) == tuple(
              bsd.output_size.scores.as_list()
          )
  
imports: import numpy
  import pytest
  import tensorflow
  from tensorflow_addons.seq2seq import attention_wrapper
  from tensorflow_addons.seq2seq import beam_search_decoder, gather_tree
label_desc: Use image node_id 16 for calling a global function with example usage: test_beam_search_decoder(cell_class, time_major, has_attention, with_alignment_history, output_all_scores) and returns: bsd, shape, unknown