code: def test_attention_state_with_keras_rnn():
      # See https://github.com/tensorflow/addons/issues/1095.
      cell = tf.keras.layers.LSTMCell(8)
  
      mechanism = wrapper.LuongAttention(
          units=8, memory=tf.ones((2, 4, 8))
      )
  
      cell = wrapper.AttentionWrapper(
          cell=cell, attention_mechanism=mechanism
      )
  
      layer = tf.keras.layers.RNN(cell)
      _ = layer(inputs=tf.ones((2, 4, 8)))
  
      # Make sure the explicit initial_state also works.
      initial_state = cell.get_initial_state(
          batch_size=2, dtype=tf.float32
      )
      _ = layer(inputs=tf.ones((2, 4, 8)), initial_state=initial_state)
  
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 25 for calling a global function with example usage: test_attention_state_with_keras_rnn() without return types