code: def test_multi_head_dot_attention_cell(
      num_heads, scaled, normalized, hybridize, rel_score_type, ctx
  ):
      with ctx:
          batch_size = 5
          query_length, mem_length = 16, 32
          query_head_units = 8
          mem_head_units = 6
          query_units = query_head_units * num_heads
          mem_units = mem_head_units * num_heads
          seed = 100
          attn_cells = dict()
          for layout in ["NKT", "NTK", "TNK"]:
              for use_einsum in [False, True]:
                  attn_cells[
                      (layout, use_einsum)
                  ] = MultiHeadAttentionCell(
                      query_units=query_units,
                      num_heads=num_heads,
                      attention_dropout=0.0,
                      scaled=scaled,
                      normalized=normalized,
                      layout=layout,
                      use_einsum=use_einsum,
                  )
                  if hybridize:
                      attn_cells[(layout, use_einsum)].hybridize()
          # Generate the data
          query_np = np.random.normal(
              0,
              1,
              (batch_size, num_heads, query_length, query_head_units),
          )
          key_np = np.random.normal(
              0,
              1,
              (batch_size, num_heads, mem_length, query_head_units),
          )
          value_np = np.random.normal(
              0, 1, (batch_size, num_heads, mem_length, mem_head_units)
          )
          mask_np = np.random.randint(
              0, 2, (batch_size, query_length, mem_length)
          )
          if rel_score_type == "share_head":
              rel_scores_np = np.random.normal(
                  0, 1, (query_length, mem_length)
              )
          elif rel_score_type == "no_share_head":
              rel_scores_np = np.random.normal(
                  0, 1, (num_heads, query_length, mem_length)
              )
          else:
              rel_scores_np = None
          out_np = None
          score_np = None
          attn_weights_np = None
          stored_layout = None
          query_grad_np = None
          key_grad_np = None
          value_grad_np = None
          rel_scores_grad_np = None
          for (layout, use_einsum), attn_cell in attn_cells.items():
              mx.npx.random.seed(seed)
              if rel_score_type != "no":
                  rel_scores = mx.np.array(
                      rel_scores_np, dtype=np.float32
                  )
              else:
                  rel_scores = None
              if layout == "NKT":
                  query = mx.np.array(query_np, dtype=np.float32)
                  key = mx.np.array(key_np, dtype=np.float32)
                  value = mx.np.array(value_np, dtype=np.float32)
              elif layout == "NTK":
                  query = mx.np.array(
                      query_np.transpose((0, 2, 1, 3)), dtype=np.float32
                  )
                  key = mx.np.array(
                      key_np.transpose((0, 2, 1, 3)), dtype=np.float32
                  )
                  value = mx.np.array(
                      value_np.transpose((0, 2, 1, 3)), dtype=np.float32
                  )
              elif layout == "TNK":
                  query = mx.np.array(
                      query_np.transpose((2, 0, 1, 3)), dtype=np.float32
                  )
                  key = mx.np.array(
                      key_np.transpose((2, 0, 1, 3)), dtype=np.float32
                  )
                  value = mx.np.array(
                      value_np.transpose((2, 0, 1, 3)), dtype=np.float32
                  )
              else:
                  raise NotImplementedError
              mask = mx.np.array(mask_np, dtype=np.int32)
              query.attach_grad()
              key.attach_grad()
              value.attach_grad()
              if rel_scores is not None:
                  rel_scores.attach_grad()
              with mx.autograd.record():
                  out, [score, attn_weights] = attn_cell(
                      query, key, value, mask, rel_scores
                  )
                  out.backward()
              if layout == "NKT":
                  assert out.shape == (
                      batch_size,
                      query_length,
                      num_heads * mem_head_units,
                  )
              elif layout == "NTK":
                  assert out.shape == (
                      batch_size,
                      query_length,
                      num_heads * mem_head_units,
                  )
              elif layout == "TNK":
                  assert out.shape == (
                      query_length,
                      batch_size,
                      num_heads * mem_head_units,
                  )
              else:
                  raise NotImplementedError
              for i in range(num_heads):
                  assert_allclose(
                      attn_weights[:, i, :, :][mask == 0].asnumpy(),
                      mask[mask == 0].astype(np.float32).asnumpy(),
                      1e-5,
                      1e-5,
                  )
  
              if stored_layout is None:
                  out_np = out.asnumpy()
                  score_np = score.asnumpy()
                  attn_weights_np = attn_weights.asnumpy()
                  stored_layout = layout
                  query_grad_np = query.grad.asnumpy()
                  key_grad_np = key.grad.asnumpy()
                  value_grad_np = value.grad.asnumpy()
                  if rel_score_type != "no":
                      rel_scores_grad_np = rel_scores.grad.asnumpy()
              else:
                  assert stored_layout == "NKT"
                  # Begin to match the output
                  if layout == "NKT":
                      m_out_np = out.asnumpy()
                      m_score_np = score.asnumpy()
                      m_attn_weights_np = attn_weights.asnumpy()
                      m_query_grad_np = query.grad.asnumpy()
                      m_key_grad_np = key.grad.asnumpy()
                      m_value_grad_np = value.grad.asnumpy()
                      if rel_score_type != "no":
                          m_rel_scores_grad_np = (
                              rel_scores.grad.asnumpy()
                          )
                  elif layout == "NTK":
                      m_out_np = out.asnumpy()
                      m_score_np = score.asnumpy()
                      m_attn_weights_np = attn_weights.asnumpy()
                      m_query_grad_np = query.grad.asnumpy().transpose(
                          (0, 2, 1, 3)
                      )
                      m_key_grad_np = key.grad.asnumpy().transpose(
                          (0, 2, 1, 3)
                      )
                      m_value_grad_np = value.grad.asnumpy().transpose(
                          (0, 2, 1, 3)
                      )
                      if rel_score_type != "no":
                          m_rel_scores_grad_np = (
                              rel_scores.grad.asnumpy()
                          )
                  elif layout == "TNK":
                      m_out_np = out.asnumpy().transpose((1, 0, 2))
                      m_score_np = score.asnumpy()
                      m_attn_weights_np = attn_weights.asnumpy()
                      m_query_grad_np = query.grad.asnumpy().transpose(
                          (1, 2, 0, 3)
                      )
                      m_key_grad_np = key.grad.asnumpy().transpose(
                          (1, 2, 0, 3)
                      )
                      m_value_grad_np = value.grad.asnumpy().transpose(
                          (1, 2, 0, 3)
                      )
                      if rel_score_type != "no":
                          m_rel_scores_grad_np = (
                              rel_scores.grad.asnumpy()
                          )
                  else:
                      raise NotImplementedError
                  assert_allclose(m_out_np, out_np, 1e-5, 1e-5)
                  assert_allclose(m_score_np, score_np, 1e-5, 1e-5)
                  assert_allclose(
                      m_attn_weights_np, attn_weights_np, 1e-5, 1e-5
                  )
                  assert_allclose(
                      m_query_grad_np, query_grad_np, 1e-5, 1e-5
                  )
                  assert_allclose(
                      m_key_grad_np, key_grad_np, 1e-5, 1e-5
                  )
                  assert_allclose(
                      m_value_grad_np, value_grad_np, 1e-5, 1e-5
                  )
                  if rel_score_type != "no":
                      assert_allclose(
                          m_rel_scores_grad_np,
                          rel_scores_grad_np,
                          1e-5,
                          1e-5,
                      )
  
file_path: gluon-nlp/tests/test_attention_cell.py
imports: import numpy
  from numpy.testing import assert_allclose
  import pytest
  import mxnet
  from mxnet.gluon import HybridBlock
  from gluonnlp.attention_cell import multi_head_dot_attn, gen_self_attn_mask, gen_mem_attn_mask, MultiHeadAttentionCell, RelAttentionScoreCell
  from gluonnlp.utils.parameter import grad_global_norm
label_desc: Use image node_id 1 for calling a global function with example usage: test_multi_head_dot_attention_cell(num_heads, scaled, normalized, hybridize, rel_score_type, ctx) without return types