code: def _test_with_attention(
      create_attention_mechanism,
      expected_final_output,
      expected_final_state,
      attention_mechanism_depth=3,
      alignment_history=False,
      expected_final_alignment_history=None,
      attention_layer_size=6,
      attention_layer=None,
      create_query_layer=False,
      create_memory_layer=True,
      create_attention_kwargs=None,
  ):
      attention_layer_sizes = (
          [attention_layer_size]
          if attention_layer_size is not None
          else None
      )
      attention_layers = (
          [attention_layer] if attention_layer is not None else None
      )
      create_attention_mechanisms = [create_attention_mechanism]
      attention_mechanism_depths = [attention_mechanism_depth]
      assert len(create_attention_mechanisms) == 1
      encoder_sequence_length = [3, 2, 3, 1, 1]
      decoder_sequence_length = [2, 0, 1, 2, 3]
      batch_size = 5
      encoder_max_time = 8
      decoder_max_time = 4
      input_depth = 7
      encoder_output_depth = 10
      cell_depth = 9
      create_attention_kwargs = create_attention_kwargs or {}
  
      if attention_layer_sizes is not None:
          # Compute sum of attention_layer_sizes. Use encoder_output_depth if
          # None.
          attention_depth = sum(
              attention_layer_size or encoder_output_depth
              for attention_layer_size in attention_layer_sizes
          )
      elif attention_layers is not None:
          # Compute sum of attention_layers output depth.
          attention_depth = sum(
              attention_layer.compute_output_shape(
                  [batch_size, cell_depth + encoder_output_depth]
              )[-1]
              for attention_layer in attention_layers
          )
      else:
          attention_depth = encoder_output_depth * len(
              create_attention_mechanisms
          )
  
      decoder_inputs = np.random.randn(
          batch_size, decoder_max_time, input_depth
      ).astype(np.float32)
      encoder_outputs = np.random.randn(
          batch_size, encoder_max_time, encoder_output_depth
      ).astype(np.float32)
  
      attention_mechanisms = []
      for creator, depth in zip(
          create_attention_mechanisms, attention_mechanism_depths
      ):
          # Create a memory layer with deterministic initializer to avoid
          # randomness in the test between graph and eager.
          if create_query_layer:
              create_attention_kwargs[
                  "query_layer"
              ] = tf.keras.layers.Dense(
                  depth, kernel_initializer="ones", use_bias=False
              )
          if create_memory_layer:
              create_attention_kwargs[
                  "memory_layer"
              ] = tf.keras.layers.Dense(
                  depth, kernel_initializer="ones", use_bias=False
              )
  
          attention_mechanisms.append(
              creator(
                  units=depth,
                  memory=encoder_outputs,
                  memory_sequence_length=encoder_sequence_length,
                  **create_attention_kwargs,
              )
          )
  
      attention_layer_size = attention_layer_sizes
      attention_layer = attention_layers
      if attention_layer_size is not None:
          attention_layer_size = attention_layer_size[0]
      if attention_layer is not None:
          attention_layer = attention_layer[0]
      cell = tf.keras.layers.LSTMCell(
          cell_depth,
          recurrent_activation="sigmoid",
          kernel_initializer="ones",
          recurrent_initializer="ones",
      )
      cell = wrapper.AttentionWrapper(
          cell,
          attention_mechanisms[0],
          attention_layer_size=attention_layer_size,
          alignment_history=alignment_history,
          attention_layer=attention_layer,
      )
      if cell._attention_layers is not None:
          for layer in cell._attention_layers:
              layer.kernel_initializer = (
                  tf.compat.v1.keras.initializers.glorot_uniform(
                      seed=1337
                  )
              )
  
      policy = tf.keras.mixed_precision.global_policy()
      sampler = sampler_py.TrainingSampler()
      my_decoder = basic_decoder.BasicDecoder(
          cell=cell, sampler=sampler
      )
      initial_state = cell.get_initial_state(
          batch_size=batch_size, dtype=policy.compute_dtype
      )
      final_outputs, final_state, _ = my_decoder(
          decoder_inputs,
          initial_state=initial_state,
          sequence_length=decoder_sequence_length,
      )
  
      assert isinstance(final_outputs, basic_decoder.BasicDecoderOutput)
      assert isinstance(final_state, wrapper.AttentionWrapperState)
  
      expected_time = max(decoder_sequence_length)
      assert (batch_size, expected_time, attention_depth) == tuple(
          final_outputs.rnn_output.shape.as_list()
      )
      assert (batch_size, expected_time) == tuple(
          final_outputs.sample_id.shape.as_list()
      )
  
      assert (batch_size, attention_depth) == tuple(
          final_state.attention.shape.as_list()
      )
      assert (batch_size, cell_depth) == tuple(
          final_state.cell_state[0].shape.as_list()
      )
      assert (batch_size, cell_depth) == tuple(
          final_state.cell_state[1].shape.as_list()
      )
  
      if alignment_history:
          state_alignment_history = (
              final_state.alignment_history.stack()
          )
          assert (expected_time, batch_size, encoder_max_time) == tuple(
              state_alignment_history.shape.as_list()
          )
          tf.nest.assert_same_structure(
              cell.state_size,
              cell.get_initial_state(
                  batch_size=batch_size, dtype=policy.compute_dtype
              ),
          )
          # Remove the history from final_state for purposes of the
          # remainder of the tests.
          final_state = final_state._replace(
              alignment_history=()
          )  # pylint: disable=protected-access
      else:
          state_alignment_history = ()
  
      final_outputs = tf.nest.map_structure(np.array, final_outputs)
      final_state = tf.nest.map_structure(np.array, final_state)
      state_alignment_history = tf.nest.map_structure(
          np.array, state_alignment_history
      )
      final_output_info = tf.nest.map_structure(
          get_result_summary, final_outputs
      )
  
      final_state_info = tf.nest.map_structure(
          get_result_summary, final_state
      )
  
      tf.nest.map_structure(
          assert_allclose_or_equal,
          expected_final_output,
          final_output_info,
      )
      tf.nest.map_structure(
          assert_allclose_or_equal,
          expected_final_state,
          final_state_info,
      )
      # by default, the wrapper emits attention as output
      if alignment_history:
          final_alignment_history_info = tf.nest.map_structure(
              get_result_summary, state_alignment_history
          )
          tf.nest.map_structure(
              assert_allclose_or_equal,
              # outputs are batch major but the stacked TensorArray is
              # time major
              expected_final_alignment_history,
              final_alignment_history_info,
          )
  
imports: import collections
  from packaging.version import Version
  import pytest
  import numpy
  import tensorflow
  from tensorflow_addons.seq2seq import attention_wrapper
  from tensorflow_addons.seq2seq import basic_decoder
  from tensorflow_addons.seq2seq import sampler
label_desc: Use image node_id 12 for calling a global function with example usage: _test_with_attention(create_attention_mechanism, expected_final_output, expected_final_state, attention_mechanism_depth, alignment_history, expected_final_alignment_history, attention_layer_size, attention_layer, create_query_layer, create_memory_layer, create_attention_kwargs) without return types