code: def test_compute_output_shape_return_attn():
      batch_size = 10
      num_heads = 8
      head_size = 12
  
      mha = MultiHeadAttention(
          head_size=head_size,
          num_heads=num_heads,
          return_attn_coef=True,
      )
  
      output_shape, attn_coef_shape = mha.compute_output_shape(
          [(batch_size, 5, 9), (batch_size, 7, 11), (batch_size, 7, 13)]
      )
  
      assert output_shape[1] == 5
      assert output_shape[2] == 13
  
      assert attn_coef_shape[0] == batch_size
      assert attn_coef_shape[1] == num_heads
      assert attn_coef_shape[2] == 5
      assert attn_coef_shape[3] == 7
  
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 8 for calling a global function with example usage: test_compute_output_shape_return_attn() without return types