code: def test_multihead_self_attn(
      output_dim: int, batch_size: int, bias: bool, use_attn_mask: bool
  ):
      seq_len = 5
      embed_dim = 8
      mha = MultiHeadAttention(
          embed_dim=embed_dim,
          num_heads=2,
          attn_dropout=0.0,
          bias=bias,
          output_dim=output_dim,
          coreml_compatible=False,
      )
      mha.eval()
  
      qkv = torch.randn(size=(seq_len, batch_size, embed_dim))
  
      attn_mask = None
      if use_attn_mask:
          attn_mask = build_attention_mask(
              context_length=seq_len,
              batch_size=batch_size,
              use_pytorch_mha=True,
          )
  
      # Pytorch MHA accepts sequence first
      out_pytorch = mha(
          x_q=qkv, use_pytorch_mha=True, attn_mask=attn_mask
      )
  
      # default works with batch-first
      qkv = qkv.transpose(0, 1)
      attn_mask = None
      if use_attn_mask:
          attn_mask = build_attention_mask(
              context_length=seq_len,
              batch_size=batch_size,
              use_pytorch_mha=False,
          )
      out_default = mha(
          x_q=qkv, use_pytorch_mha=False, attn_mask=attn_mask
      )
      out_default = out_default.transpose(0, 1)
  
      torch.testing.assert_close(
          actual=out_default, expected=out_pytorch, atol=1e-3, rtol=1e-3
      )
  
      if hasattr(mha, "forward_tracing") and attn_mask is None:
          # check coreml compatible version
          out_tracing = mha.forward_tracing(x_q=qkv).transpose(0, 1)
          torch.testing.assert_close(
              actual=out_default,
              expected=out_tracing,
              atol=1e-3,
              rtol=1e-3,
          )
  
imports: import sys
  import numpy
  import pytest
  import torch
  from cvnets.layers.multi_head_attention import MultiHeadAttention
label_desc: Use image node_id 2 for calling a global function with example usage: test_multihead_self_attn(output_dim, batch_size, bias, use_attn_mask) without return types