code: def test_softmax(BLOCK, WIDTH, DTYPE):
      # set seed
      torch.random.manual_seed(0)
      Z, H, M, N = 2, 4, WIDTH, WIDTH
      scale = 0.4
  
      # create inputs
      layout = torch.randint(2, (H, M // BLOCK, N // BLOCK))
      x = torch.randn(
          (Z, H, M, N), dtype=DTYPE, requires_grad=True, device="cuda"
      )
  
      # triton result
      op = blocksparse_softmax(
          layout, BLOCK, device=torch.device("cuda")
      )
      tx = block_sparsify_tensor(x, layout, BLOCK)
      ty = op(tx, scale=scale)
  
      # torch result
      rx = mask_tensor(x, layout, BLOCK, value=float("-inf"))
      rx = rx[:, :, : (M // BLOCK) * BLOCK, : (M // BLOCK) * BLOCK]
  
      ry = torch.softmax(rx * scale, -1)
      ry = block_sparsify_tensor(ry, layout, BLOCK)
  
      # compare
      torch.testing.assert_close(ry, ty)
  
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 4 for calling a global function with example usage: test_softmax(BLOCK, WIDTH, DTYPE) without return types