code: def test_manual_memory_reset(attention_cls):
      dummy_data = DummyData()
      attention = attention_cls(dummy_data.units)
  
      def _compute_score(batch_size=None):
          if batch_size is None:
              batch_size = dummy_data.batch
          memory = dummy_data.memory[:batch_size]
          attention.setup_memory(
              memory,
              memory_sequence_length=dummy_data.memory_length[
                  :batch_size
              ],
          )
          assert attention.values.shape.as_list() == list(memory.shape)
          assert attention.keys.shape.as_list() == list(memory.shape)[
              :-1
          ] + [dummy_data.units]
          return attention(
              [
                  dummy_data.query[:batch_size],
                  dummy_data.state[:batch_size],
              ]
          )
  
      _compute_score(batch_size=dummy_data.batch)
      variables = list(attention.variables)
      _compute_score(batch_size=dummy_data.batch - 1)
  
      # No new variables were created.
      for var_1, var_2 in zip(variables, list(attention.variables)):
          assert var_1 is var_2
  
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 6 for calling a global function with example usage: test_manual_memory_reset(attention_cls) and returns: attention