code: def test_add_bias_parity():
      # values don't matter for this test.
      mha = MultiheadAttention(
          embed_dim=8,
          num_heads=2,
          dropout=0.0,
          add_bias_kv=True,
          add_zero_attn=True,
      )
  
      def old_bias_code(k, v, key_padding_mask, attn_mask, bsz):
          k = torch.cat([k, mha.bias_k.repeat(1, bsz, 1)])
          v = torch.cat([v, mha.bias_v.repeat(1, bsz, 1)])
          if attn_mask is not None:
              attn_mask = torch.cat(
                  [
                      attn_mask,
                      attn_mask.new_zeros(attn_mask.size(0), 1),
                  ],
                  dim=1,
              )
          if key_padding_mask is not None:
              key_padding_mask = torch.cat(
                  [
                      key_padding_mask,
                      key_padding_mask.new_zeros(
                          key_padding_mask.size(0), 1
                      ),
                  ],
                  dim=1,
              )
          return k, v, key_padding_mask, attn_mask
  
      seq_len = 64
      bsz = 8
      embedding = 8
      key_padding_mask = torch.rand((bsz, seq_len))
      attn_mask = torch.rand((seq_len, seq_len))
      k = torch.rand((seq_len, bsz, embedding))
      v = torch.rand((seq_len, bsz, embedding))
  
      k_orig, v_orig, kp_mask_orig, a_mask_orig = old_bias_code(
          k, v, key_padding_mask, attn_mask, bsz
      )
      k_new, v_new, kp_mask_new, a_mask_new = mha._add_bias(
          k, v, key_padding_mask, attn_mask, bsz
      )
  
      assert torch.equal(k_orig, k_new)
      assert torch.equal(v_orig, v_new)
      assert torch.equal(kp_mask_orig, kp_mask_new)
      assert torch.equal(a_mask_orig, a_mask_new)
  
file_path: fairseq/tests/test_multihead_attention.py
imports: import random
  import unittest
  import pytest
  import torch
  from fairseq.modules.multihead_attention import MultiheadAttention, _mask_for_xformers
label_desc: Use image node_id 8 for calling a global function with example usage: test_add_bias_parity() and returns: k,  v,  key_padding_mask,  attn_mask