code: def test_attention_fwd_bwd(
      block,
      dtype,
      input_scale=1.0,
      scale=1 / 8.0,
      n_ctx=384,
      batch_size=2,
      n_heads=2,
  ):
      # inputs
      head_dim = 64
      qkv_shape = (batch_size, n_heads, n_ctx, head_dim)
      qkvs = [
          torch.nn.Parameter(
              input_scale * torch.randn(qkv_shape), requires_grad=True
          )
          .to(dtype)
          .cuda()
          for _ in range(3)
      ]
  
      def loss_fn(x):
          return (x**2).mean()
  
      # Triton:
      n_blocks = n_ctx // block
      layout = torch.ones(
          [n_heads, n_blocks, n_blocks], dtype=torch.long
      )
      query, key, value = [x.clone() for x in qkvs]
      query.retain_grad()
      key.retain_grad()
      value.retain_grad()
      if block not in [16, 32, 64, 128]:
          # Check that unsupported dimensions are caught
          with pytest.raises(AssertionError):
              _ = BlockSparseAttention(layout, block)
      else:
          block_sparse_attention = BlockSparseAttention(layout, block)
          attn_out = block_sparse_attention(
              q=query, k=key, v=value, scale=scale
          )
  
          # ad hoc loss
          loss = loss_fn(attn_out)
          loss.backward()
          grads = [query.grad, key.grad, value.grad]
  
          # Torch version:
          torch_q, torch_k, torch_v = [x.clone() for x in qkvs]
          torch_q = torch_q * scale
          torch_q.retain_grad()
          torch_k.retain_grad()
          torch_v.retain_grad()
          scores = scale * torch.einsum(
              "bhsd,bhtd->bhst", torch_q, torch_k
          )
          probs = torch.softmax(scores, dim=-1)
          torch_attn_out = torch.einsum(
              "bhst,bhtd->bhsd", probs, torch_v
          )
  
          # ad hoc loss
          torch_loss = loss_fn(torch_attn_out)
          torch_loss.backward()
          torch_grads = [torch_q.grad, torch_k.grad, torch_v.grad]
  
          # comparison
          torch.testing.assert_close(
              loss,
              torch_loss,
              msg=f"Triton loss {loss} and torch loss {torch_loss}",
          )
  
          for g1, g2 in zip(grads, torch_grads):
              torch.testing.assert_close(
                  torch.norm(g1),
                  torch.norm(g2),
                  msg=f"Triton grad {torch.norm(g1).item()} and torch grad {torch.norm(g2).item()}",
              )
  
file_path: xformers/tests/test_triton_blocksparse.py
imports: import functools
  import pytest
  import torch
  import xformers
  from xformers.components import MultiHeadDispatch
  from xformers.components.attention import build_attention
  from xformers.components.attention.attention_patterns import block_sparsify_tensor
label_desc: Use image node_id 5 for calling a global function with example usage: test_attention_fwd_bwd(block, dtype, input_scale, scale, n_ctx, batch_size, n_heads) and returns: unknown