code: def test_multihead_attention():
      for query_units, key_units, value_units, num_heads in [
          (4, 4, 8, 2),
          (3, 3, 9, 3),
          (6, 6, 5, 1),
      ]:
          for use_mask in [True, False]:
              for scaled in [True, False]:
                  for normalized in [True, False]:
                      for layout in ["NKT", "NTK", "TNK"]:
                          cell = ac.MultiHeadAttentionCell(
                              query_units=query_units,
                              num_heads=num_heads,
                              scaled=scaled,
                              normalized=normalized,
                              layout=layout,
                          )
                          check_attention_cell_basic(
                              cell,
                              q_channel=query_units // num_heads,
                              k_channel=key_units // num_heads,
                              v_channel=value_units // num_heads,
                              use_mask=use_mask,
                              multi_head=True,
                              num_heads=num_heads,
                              layout=layout,
                          )
  
imports: import numpy
  from numpy.testing import assert_allclose
  import torch
  from gluonnlp.torch import attention_cell
label_desc: Use image node_id 2 for calling a global function with example usage: test_multihead_attention() without return types