code: def test_memory_re_setup(attention_cls):
      class MyModel(tf.keras.models.Model):
          def __init__(self, vocab, embedding_dim, memory_size, units):
              super().__init__()
              self.emb = tf.keras.layers.Embedding(
                  vocab, embedding_dim, mask_zero=True
              )
              self.encoder = tf.keras.layers.LSTM(
                  memory_size, return_sequences=True
              )
              self.attn_mch = attention_cls(units)
  
          def call(self, inputs):
              enc_input, query, state = inputs
              mask = self.emb.compute_mask(enc_input)
              enc_input = self.emb(enc_input)
              enc_output = self.encoder(enc_input, mask=mask)
              # To ensure manual resetting also works in the graph mode,
              # we call the attention mechanism twice.
              self.attn_mch(enc_output, mask=mask, setup_memory=True)
              self.attn_mch(enc_output, mask=mask, setup_memory=True)
              score = self.attn_mch([query, state])
              return score
  
      vocab = 20
      embedding_dim = 6
      num_batches = 5
  
      dummy_data = DummyData()
      model = MyModel(
          vocab, embedding_dim, dummy_data.memory_size, dummy_data.units
      )
      model.compile("rmsprop", "mse")
  
      x = np.random.randint(
          vocab,
          size=(num_batches * dummy_data.batch, dummy_data.timestep),
      )
      x_test = np.random.randint(
          vocab,
          size=(num_batches * dummy_data.batch, dummy_data.timestep),
      )
      y = np.random.randn(
          num_batches * dummy_data.batch, dummy_data.timestep
      )
  
      query = np.tile(dummy_data.query, [num_batches, 1])
      state = np.tile(dummy_data.state, [num_batches, 1])
  
      model.fit([x, query, state], (y, y), batch_size=dummy_data.batch)
      model.predict_on_batch([x_test, query, 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 8 for calling a global function with example usage: test_memory_re_setup(attention_cls) and returns: score