code: def test_xformers_blocksparse_parity(
      device,
      add_zero_attn,
      batch_size,
      embedding,
      seq_len,
      num_heads,
  ):
      xformers_att_config = '{"name": "scaled_dot_product"}'
      xformers_blocksparse_blocksize = 16
      xformers_blocksparse_layout = torch.ones(
          seq_len // xformers_blocksparse_blocksize,
          seq_len // xformers_blocksparse_blocksize,
          dtype=torch.int32,
      )
  
      q = torch.rand(seq_len, batch_size, embedding).to(device).half()
      q.requires_grad = True
      k = torch.rand(seq_len, batch_size, embedding).to(device).half()
      k.requires_grad = True
      v = torch.rand(seq_len, batch_size, embedding).to(device).half()
      v.requires_grad = True
  
      q_ = q.detach().clone().half()
      q_.requires_grad = True
      k_ = k.detach().clone().half()
      k_.requires_grad = True
      v_ = v.detach().clone().half()
      v_.requires_grad = True
  
      _reset_seeds()
      xf_blocksparse_mha = (
          MultiheadAttention(
              embedding,
              num_heads,
              dropout=0.0,
              add_zero_attn=add_zero_attn,
              xformers_att_config=xformers_att_config,
              xformers_blocksparse_layout=xformers_blocksparse_layout,
              xformers_blocksparse_blocksize=xformers_blocksparse_blocksize,
          )
          .to(device)
          .half()
      )
  
      xf_blocksparse_output, _ = xf_blocksparse_mha(
          q,
          k,
          v,
      )
  
      _reset_seeds()
      xformers_mha = (
          MultiheadAttention(
              embedding,
              num_heads,
              dropout=0.0,
              add_zero_attn=add_zero_attn,
              xformers_att_config=xformers_att_config,
              xformers_blocksparse_layout=None,
          )
          .to(device)
          .half()
      )
  
      xformers_output, _ = xformers_mha(
          q_,
          k_,
          v_,
      )
  
      # # account for when nan != nan
      rand = random.uniform(0, 1)
      xformers_output = xformers_output.masked_fill(
          xformers_output.isnan(), rand
      )
      xf_blocksparse_output = xf_blocksparse_output.masked_fill(
          xf_blocksparse_output.isnan(), rand
      )
  
      assert_almost_equal(xformers_output, xf_blocksparse_output)
  
      loss_blocksparse = torch.norm(xformers_output)
      loss_original = torch.norm(xf_blocksparse_output)
      loss_blocksparse.backward()
      loss_original.backward()
  
      q.masked_fill(q.isnan(), rand)
      q_.masked_fill(q_.isnan(), rand)
      k.masked_fill(k.isnan(), rand)
      k_.masked_fill(k_.isnan(), rand)
      v.masked_fill(v.isnan(), rand)
      v_.masked_fill(v_.isnan(), rand)
  
      assert_almost_equal(q.grad, q_.grad)
      assert_almost_equal(k.grad, k_.grad)
      assert_almost_equal(v.grad, v_.grad)
  
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 5 for calling a global function with example usage: test_xformers_blocksparse_parity(device, add_zero_attn, batch_size, embedding, seq_len, num_heads) without return types