code: def test_multihead_attention(self, average_attn_weights):
      def _scaled_dot_attn_ref(
          Q,
          K,
          V,
          dims,
          unseen_mask=None,
          key_padding_mask=None,
          average_attn_weights=average_attn_weights,
      ):
          """Numpy-based reference implementation of scaled dot attention
          for testing"""
  
          QKT = _batchmatmul(
              Q,
              np.transpose(K, axes=[0, 1, 3, 2])
              / np.sqrt(
                  dims[3], dtype=np.float32
              ),  # divide by sqrt(d_head)
          )
          b1, b2, s1, s2 = QKT.shape
          if unseen_mask is not None or key_padding_mask is not None:
              # assert s1 == s2
              for i in range(b1):
                  for j in range(b2):
                      for m in range(s1):
                          for n in range(s2):
                              if (
                                  unseen_mask is not None
                                  and unseen_mask[m][n] == 0
                              ):
                                  QKT[i, j, m, n] = -np.inf
                              if (
                                  key_padding_mask is not None
                                  and key_padding_mask[i][n]
                              ):
                                  QKT[i, j, m, n] = -np.inf
  
          reference = _softmax(QKT)
          ref_attn_weight = reference
          if average_attn_weights:
              ref_attn_weight = np.sum(ref_attn_weight, axis=1) / b2
          reference = _batchmatmul(reference, V)
          return reference, ref_attn_weight
  
      def _batchmatmul(a, b):  # batchmatmul over 4 dim matrix
          """Numpy-based batch matrix multiply over 4 dim matrix"""
          assert a.shape[0] == b.shape[0]
          assert a.shape[1] == b.shape[1]
          retval = np.zeros(
              (a.shape[0], a.shape[1], a.shape[2], b.shape[3]),
              dtype=np.float32,
          )
          for i in range(a.shape[0]):
              for j in range(a.shape[1]):
                  retval[i, j, :, :] = np.matmul(
                      a[i, j, :, :], b[i, j, :, :]
                  )
          return retval
  
      def _softmax(x):  # softmax over 4 dim matrix
          """Numpy-based reference softmax over 4 dim matrix"""
          np.seterr(invalid="ignore")
          output = np.zeros(x.shape, dtype=np.float64)
          for i in range(x.shape[0]):
              for j in range(x.shape[1]):
                  for k in range(x.shape[2]):
                      x_curr = x[i, j, k, :]
                      e_x = np.exp(x_curr - np.amax(x_curr))
                      output[i, j, k, :] = e_x / np.sum(e_x)
          return output
  
      def _split_heads_ref(X, dims, nheads, d_head):
          X_split = np.reshape(X, dims[:2] + [nheads, d_head])
          X_split_transposed = np.transpose(X_split, [0, 2, 1, 3])
          reference = np.reshape(
              X_split_transposed, [dims[0], nheads, dims[1], d_head]
          )
          return reference
  
      def _combine_heads_ref(X, dims, nheads, d_head):
          X_transposed = np.transpose(X, [0, 2, 1, 3])
          reference = np.reshape(
              X_transposed, dims[:2] + [nheads * d_head]
          )
          return reference
  
      def _fc(X, X_weight, X_bias):
          X_fc_b = X_bias.detach().numpy()
          X_fc_w = X_weight.detach().numpy()
          return np.matmul(X, np.transpose(X_fc_w)) + X_fc_b
  
      def _create_src_lengths_mask(batch_size, src_lengths):
          """
          Generate boolean mask to prevent attention beyond the end of source
          Inputs:
            batch_size : int
            src_lengths : [batch_size] of sentence lengths
          Outputs:
            [batch_size, max_src_len]
          """
          max_srclen = src_lengths.max()
          src_indices = (
              torch.arange(0, max_srclen).unsqueeze(0).to(src_lengths)
          )
          src_indices = src_indices.expand(batch_size, max_srclen)
          src_lengths = src_lengths.unsqueeze(dim=1).expand(
              batch_size, max_srclen
          )
          # returns [batch_size, max_seq_len]
          return (src_indices < src_lengths).int().detach()
  
      def _multihead_attn_test_helper(
          add_key_padding_mask=False,
          add_bias_kv=False,
          add_zero_attn=False,
          saved_kv=False,
          same_embed_dim=False,
          average_attn_weights=average_attn_weights,
      ):
          for _ in range(100):
              batch_sz, seq_len = (
                  random.randint(2, 10) for r in range(2)
              )
              d_head = random.randint(3, 10)
              nheads = random.randint(2, 5) * 2
              d_model = d_head * nheads
              if same_embed_dim:
                  kv_dim = d_model
              else:
                  kv_dim = random.randint(5, 20)
              dims = [batch_sz, seq_len, kv_dim]
  
              saved_k = None
              saved_k_tensor = None
              saved_v = None
              saved_v_tensor = None
              if saved_kv:
                  saved_k = np.random.rand(
                      batch_sz * nheads, seq_len, d_head
                  )
                  saved_k_tensor = torch.from_numpy(saved_k).to(
                      torch.get_default_dtype()
                  )
                  saved_v = np.random.rand(
                      batch_sz * nheads, seq_len, d_head
                  )
                  saved_v_tensor = torch.from_numpy(saved_v).to(
                      torch.get_default_dtype()
                  )
  
              key_padding_mask = None
              key_padding_mask_tensor = None
              if add_key_padding_mask:
                  seq_mask = np.random.randint(0, 2, (1, seq_len))
                  key_padding_mask = (
                      np.repeat(seq_mask, batch_sz, axis=0) == 1
                  )
                  key_padding_mask_tensor = torch.from_numpy(
                      key_padding_mask
                  )
              decoder_state = np.random.rand(batch_sz, d_model)
              K = np.random.rand(*dims)
              V = K
              Q = np.expand_dims(decoder_state, 1)
              attn_mask = np.random.randint(0, 2, size=(1, seq_len))
              attn_mask_tensor = torch.from_numpy(attn_mask).float()
              attn_mask_tensor.masked_fill_(
                  attn_mask_tensor == 0, float("-inf")
              )
              attn_mask_tensor.masked_fill_(
                  attn_mask_tensor > 0, float("0.0")
              )
  
              decoder_state_tensor = torch.from_numpy(decoder_state).to(
                  torch.get_default_dtype()
              )
              source_hid_tensor = (
                  torch.from_numpy(K)
                  .to(torch.get_default_dtype())
                  .transpose(0, 1)
              )
  
              multihead_attn_module = MultiheadAttention(
                  d_model,
                  nheads,
                  add_bias_kv=add_bias_kv,
                  add_zero_attn=add_zero_attn,
                  kdim=kv_dim,
                  vdim=kv_dim,
              )
  
              if add_bias_kv:
                  bias_k = multihead_attn_module.bias_k.detach().numpy()
                  bias_v = multihead_attn_module.bias_v.detach().numpy()
              else:
                  bias_k = None
                  bias_v = None
  
              _Q = decoder_state_tensor.unsqueeze(1).transpose(0, 1)
              _V = source_hid_tensor
              _K = source_hid_tensor
  
              if multihead_attn_module._qkv_same_embed_dim:
                  (
                      result,
                      result_weight,
                  ) = torch.nn.functional.multi_head_attention_forward(
                      _Q,
                      _K,
                      _V,
                      d_model,
                      nheads,
                      multihead_attn_module.in_proj_weight,
                      multihead_attn_module.in_proj_bias,
                      multihead_attn_module.bias_k,
                      multihead_attn_module.bias_v,
                      multihead_attn_module.add_zero_attn,
                      multihead_attn_module.dropout,
                      multihead_attn_module.out_proj.weight,
                      multihead_attn_module.out_proj.bias,
                      multihead_attn_module.training,
                      key_padding_mask_tensor,
                      True,
                      attn_mask_tensor,
                      static_k=saved_k_tensor,
                      static_v=saved_v_tensor,
                      average_attn_weights=average_attn_weights,
                      is_causal=False,
                  )
              else:
                  (
                      result,
                      result_weight,
                  ) = torch.nn.functional.multi_head_attention_forward(
                      _Q,
                      _K,
                      _V,
                      d_model,
                      nheads,
                      None,
                      multihead_attn_module.in_proj_bias,
                      multihead_attn_module.bias_k,
                      multihead_attn_module.bias_v,
                      multihead_attn_module.add_zero_attn,
                      multihead_attn_module.dropout,
                      multihead_attn_module.out_proj.weight,
                      multihead_attn_module.out_proj.bias,
                      multihead_attn_module.training,
                      key_padding_mask_tensor,
                      True,
                      attn_mask_tensor,
                      True,
                      multihead_attn_module.q_proj_weight,
                      multihead_attn_module.k_proj_weight,
                      multihead_attn_module.v_proj_weight,
                      static_k=saved_k_tensor,
                      static_v=saved_v_tensor,
                      average_attn_weights=average_attn_weights,
                      is_causal=False,
                  )
  
              result = result.squeeze(0).detach().numpy()
  
              if multihead_attn_module._qkv_same_embed_dim:
                  q_proj_weight = multihead_attn_module.in_proj_weight[
                      :d_model
                  ]
                  k_proj_weight = multihead_attn_module.in_proj_weight[
                      d_model : (d_model * 2)
                  ]
                  v_proj_weight = multihead_attn_module.in_proj_weight[
                      (d_model * 2) :
                  ]
              else:
                  q_proj_weight = multihead_attn_module.q_proj_weight
                  k_proj_weight = multihead_attn_module.k_proj_weight
                  v_proj_weight = multihead_attn_module.v_proj_weight
  
              Q_fc = _fc(
                  Q,
                  q_proj_weight,
                  multihead_attn_module.in_proj_bias[:d_model],
              )
              K_fc = _fc(
                  K,
                  k_proj_weight,
                  multihead_attn_module.in_proj_bias[
                      d_model : (d_model * 2)
                  ],
              )
              V_fc = _fc(
                  V,
                  v_proj_weight,
                  multihead_attn_module.in_proj_bias[(d_model * 2) :],
              )
  
              if add_bias_kv:
                  K_fc = np.concatenate(
                      (K_fc, np.repeat(bias_k, K_fc.shape[0], axis=0)),
                      axis=1,
                  )
                  V_fc = np.concatenate(
                      (V_fc, np.repeat(bias_v, V_fc.shape[0], axis=0)),
                      axis=1,
                  )
                  if attn_mask is not None:
                      attn_mask = np.concatenate(
                          (attn_mask, np.ones([1, 1])), axis=1
                      )
                  if key_padding_mask is not None:
                      key_padding_mask = np.concatenate(
                          (
                              key_padding_mask,
                              np.full((batch_sz, 1), False, dtype=bool),
                          ),
                          axis=1,
                      )
                  dims[1] += 1
              Q_split = _split_heads_ref(
                  Q_fc, [batch_sz, 1, d_model], nheads, d_head
              )
  
              if saved_k is not None:
                  K_split = np.reshape(
                      saved_k, [dims[0], nheads, dims[1], d_head]
                  )
              else:
                  K_split = _split_heads_ref(K_fc, dims, nheads, d_head)
  
              if saved_v is not None:
                  V_split = np.reshape(
                      saved_v, [dims[0], nheads, dims[1], d_head]
                  )
              else:
                  V_split = _split_heads_ref(V_fc, dims, nheads, d_head)
  
              if add_zero_attn:
                  dims[1] += 1
                  K_split = np.concatenate(
                      (
                          K_split,
                          np.zeros(
                              [
                                  K_split.shape[0],
                                  K_split.shape[1],
                                  1,
                                  K_split.shape[3],
                              ]
                          ),
                      ),
                      axis=2,
                  )
                  V_split = np.concatenate(
                      (
                          V_split,
                          np.zeros(
                              [
                                  V_split.shape[0],
                                  V_split.shape[1],
                                  1,
                                  V_split.shape[3],
                              ]
                          ),
                      ),
                      axis=2,
                  )
  
                  if attn_mask is not None:
                      attn_mask = np.concatenate(
                          (attn_mask, np.ones([1, 1])), axis=1
                      )
  
                  if key_padding_mask is not None:
                      key_padding_mask = np.concatenate(
                          (
                              key_padding_mask,
                              np.full((batch_sz, 1), False, dtype=bool),
                          ),
                          axis=1,
                      )
              attn_heads, ref_attn_weight = _scaled_dot_attn_ref(
                  Q=Q_split,
                  K=K_split,
                  V=V_split,
                  dims=Q_split.shape,
                  unseen_mask=attn_mask,
                  key_padding_mask=key_padding_mask,
              )
              combined_attn_heads = _combine_heads_ref(
                  X=attn_heads,
                  dims=[batch_sz, 1],
                  nheads=nheads,
                  d_head=d_head,
              )
  
              reference = _fc(
                  combined_attn_heads,
                  multihead_attn_module.out_proj.weight,
                  multihead_attn_module.out_proj.bias,
              )
              reference = np.squeeze(reference, axis=1)
  
              # result = reference
              self.assertEqual(tuple(result.shape), (batch_sz, d_model))
              np.testing.assert_allclose(result, reference, atol=1e-5)
  
              # result_weight = ref_attn_weight
              result_weight = result_weight.detach().numpy()
              self.assertEqual(
                  tuple(result_weight.shape),
                  tuple(ref_attn_weight.shape),
              )
              np.testing.assert_allclose(
                  result_weight, ref_attn_weight, atol=1e-5
              )
  
      def test_multihead_attn_add_bias_kv():
          _multihead_attn_test_helper(add_bias_kv=True)
  
      def test_multihead_attn_add_zero_attn():
          _multihead_attn_test_helper(add_zero_attn=True)
  
      def test_multihead_attn_no_masking():
          _multihead_attn_test_helper()
  
      def test_multihead_attn_key_padding_mask():
          _multihead_attn_test_helper(add_key_padding_mask=True)
  
      def test_multihead_attn_saved_kv():
          _multihead_attn_test_helper(saved_kv=True)
  
      def test_multihead_attn_add_bias_kv_zero_attn():
          _multihead_attn_test_helper(
              add_key_padding_mask=True,
              add_bias_kv=True,
              add_zero_attn=True,
          )
  
      def test_multihead_attn_all_arguments1():
          _multihead_attn_test_helper(
              add_key_padding_mask=True,
              add_zero_attn=True,
              saved_kv=True,
          )
  
      def test_multihead_attn_all_arguments2():
          _multihead_attn_test_helper(
              add_key_padding_mask=True,
              add_bias_kv=True,
              add_zero_attn=True,
              saved_kv=True,
          )
  
      def test_multihead_attn_all_arguments3():
          _multihead_attn_test_helper(
              add_key_padding_mask=True,
              add_zero_attn=True,
              saved_kv=True,
              same_embed_dim=True,
          )
  
      test_multihead_attn_add_zero_attn()  # Test MultiheadAttention with add_zero_attn
      test_multihead_attn_add_bias_kv()  # Test MultiheadAttention with add_bias_kv
      test_multihead_attn_no_masking()  # Test MultiheadAttention without masking
      test_multihead_attn_key_padding_mask()  # Test MultiheadAttention with src lengths
      test_multihead_attn_saved_kv()  # Test MultiheadAttention with static kv.
      test_multihead_attn_add_bias_kv_zero_attn()  # Test MultiheadAttention with bias_kv and zero_attn.
      test_multihead_attn_all_arguments1()  # Test MultiheadAttention with all the argument.
      with self.assertRaisesRegex(
          AssertionError, "bias cannot be added to static key."
      ):
          test_multihead_attn_all_arguments2()  # Test MultiheadAttention with all the argument.
      test_multihead_attn_all_arguments3()
  
imports: import contextlib
  import random
  import unittest
  import unittest.mock
  from torch.nn import MultiheadAttention
  from torch.testing._internal.common_device_type import instantiate_device_type_tests, dtypes, onlyCUDAAndPRIVATEUSE1
  from torch.testing._internal.common_nn import NNTestCase
  from torch.testing._internal.common_utils import run_tests, TEST_NUMPY, TEST_WITH_CROSSREF, parametrize, instantiate_parametrized_tests
  import torch.nn
  import torch
label_desc: Use image node_id 1 for calling the TestMultiheadAttentionNN obj's underlying member method code with example usage: obj.test_multihead_attention(average_attn_weights) and returns: reference,  ref_attn_weight, retval, output, reference, reference, unknown, unknown