code: def test_extra_dims():
      batch_size = 10
      extra_dim = 17
      num_heads = 8
      head_size = 12
  
      q = tf.random.uniform(
          (batch_size, extra_dim, 5, 9), dtype=np.float32
      )
      k = tf.random.uniform(
          (batch_size, extra_dim, 7, 11), dtype=np.float32
      )
      v = tf.random.uniform(
          (batch_size, extra_dim, 7, 13), dtype=np.float32
      )
  
      mha = MultiHeadAttention(head_size=head_size, num_heads=num_heads)
  
      output = mha([q, k, v])
  
      assert output.shape[0] == batch_size
      assert output.shape[1] == extra_dim
      assert output.shape[2] == q.shape[2]
      assert output.shape[3] == v.shape[3]
  
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 4 for calling a global function with example usage: test_extra_dims() without return types