code: def test_save_load_model():
      num_heads = 8
      head_size = 12
  
      inputs = tf.keras.layers.Input(shape=[42, 13])
  
      net, attn_coef = MultiHeadAttention(
          head_size=head_size,
          num_heads=num_heads,
          dropout=0.5,
          return_attn_coef=True,
      )([inputs, inputs, inputs])
      net = tf.keras.layers.GlobalAveragePooling1D()(net)
      net = tf.keras.layers.Dense(10, activation="softmax")(net)
  
      model = tf.keras.Model(inputs=inputs, outputs=[net, attn_coef])
  
      # initialize model
      model.predict(np.random.uniform(size=(10, 42, 13)))
  
      with tempfile.TemporaryDirectory() as model_dir:
          model_path = str(Path(model_dir) / "saved_model")
          model.save(model_path)
          new_model = tf.keras.models.load_model(model_path)
  
      assert (
          model.layers[1].get_config()
          == new_model.layers[1].get_config()
      )
  
file_path: addons/tensorflow_addons/layers/tests/multihead_attention_test.py
imports: import tempfile
  from pathlib import Path
  import numpy
  import pytest
  import tensorflow
  from tensorflow_addons.layers.multihead_attention import MultiHeadAttention
label_desc: Use image node_id 12 for calling a global function with example usage: test_save_load_model() without return types