code: def test_save_load_layer(attention_cls):
      dummy_data = DummyData()
      vocab = 20
      embedding_dim = 6
      inputs = tf.keras.Input(shape=[dummy_data.timestep])
      encoder_input = tf.keras.layers.Embedding(
          vocab, embedding_dim, mask_zero=True
      )(inputs)
      encoder_output = tf.keras.layers.LSTM(
          dummy_data.memory_size, return_sequences=True
      )(encoder_input)
  
      attention = attention_cls(dummy_data.units, encoder_output)
      query = tf.keras.Input(shape=[dummy_data.units])
      state = tf.keras.Input(shape=[dummy_data.timestep])
  
      score = attention([query, state])
  
      x_test = np.random.randint(
          vocab, size=(dummy_data.batch, dummy_data.timestep)
      )
      model = tf.keras.Model([inputs, query, state], score)
      # Fall back to v1 style Keras training loop until issue with
      # using outputs of a layer in another layer's constructor.
      model.compile("rmsprop", "mse")
      y_ref = model.predict_on_batch(
          [x_test, dummy_data.query, dummy_data.state]
      )
  
      if Version(tf.__version__) >= Version("2.13"):
          model.use_legacy_config = True
  
      config = model.get_config()
      weights = model.get_weights()
      loaded_model = tf.keras.Model.from_config(
          config, custom_objects={attention_cls.__name__: attention_cls}
      )
      loaded_model.set_weights(weights)
  
      # Fall back to v1 style Keras training loop until issue with
      # using outputs of a layer in another layer's constructor.
      loaded_model.compile("rmsprop", "mse")
  
      y = loaded_model.predict_on_batch(
          [x_test, dummy_data.query, dummy_data.state]
      )
  
      np.testing.assert_allclose(y_ref, y)
  
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 5 for calling a global function with example usage: test_save_load_layer(attention_cls) without return types