code: def test_softmax(
      BLOCK, WIDTH, is_dense, Z=2, H=2, is_causal=True, scale=0.4
  ):
      # set seed
      torch.random.manual_seed(0)
      Z, H, M, N = 2, 3, WIDTH, WIDTH
      # initialize layout
      # make sure each row has at least one non-zero element
      layout = torch.randint(2, (H, M // BLOCK, N // BLOCK))
      if is_dense:
          layout[:] = 1
      else:
          layout[1, 2, :] = 0
          layout[1, :, 1] = 0
      # initialize data
      a_shape = (Z, H, M, N)
      a_ref, a_tri = make_pair(a_shape)
      dout_ref, dout_tri = make_pair(a_shape)
      # compute [torch]
      a_ref = mask_tensor(a_ref, layout, BLOCK, value=float("-inf"))
      a_ref.retain_grad()
      at_mask = torch.ones((M, N), device="cuda")
      if is_causal:
          at_mask = torch.tril(at_mask)
      M = at_mask[None, None, :, :] + torch.zeros_like(a_ref)
      a_ref[M == 0] = float("-inf")
      out_ref = torch.softmax(a_ref * scale, -1)
      out_ref.backward(dout_ref)
      out_ref = sparsify_tensor(out_ref, layout, BLOCK)
      da_ref = sparsify_tensor(a_ref.grad, layout, BLOCK)
      # compute [triton]
      a_tri = sparsify_tensor(a_tri, layout, BLOCK)
      a_tri.retain_grad()
      dout_tri = sparsify_tensor(dout_tri, layout, BLOCK)
      op = triton.ops.blocksparse.softmax(
          layout, BLOCK, device="cuda", is_dense=is_dense
      )
      out_tri = op(a_tri, scale=scale, is_causal=is_causal)
      out_tri.backward(dout_tri)
      da_tri = a_tri.grad
      # compare
      torch.testing.assert_close(out_tri, out_ref, equal_nan=True)
      torch.testing.assert_close(da_tri, da_ref, equal_nan=True)
  
file_path: triton/python/test/unit/operators/test_blocksparse.py
imports: import pytest
  import torch
  import triton
  import triton.ops
label_desc: Use image node_id 5 for calling a global function with example usage: test_softmax(BLOCK, WIDTH, is_dense, Z, H, is_causal, scale) without return types