code: def test_multihead_cross_attn(
      output_dim: int, key_len: int, batch_size: int, bias: bool
  ):
      seq_len = 20
      embed_dim = 32
      mha = MultiHeadAttention(
          embed_dim=embed_dim,
          num_heads=8,
          attn_dropout=0.0,
          bias=bias,
          output_dim=output_dim,
          coreml_compatible=False,
      )
      mha.eval()
  
      query = torch.randn(size=(seq_len, batch_size, embed_dim))
      key = torch.randn(size=(key_len, batch_size, embed_dim))
  
      # Pytorch MHA accepts sequence first
      out_pytorch = mha(x_q=query, x_kv=key, use_pytorch_mha=True)
  
      # default works with batch-first
      query = query.transpose(0, 1)
      key = key.transpose(0, 1)
      out_default = mha(x_q=query, x_kv=key)
      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"):
          # check coreml compatible version
          out_tracing = mha.forward_tracing(
              x_q=query, x_kv=key
          ).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 3 for calling a global function with example usage: test_multihead_cross_attn(output_dim, key_len, batch_size, bias) without return types