code: def test_softmax(block, width, dtype):
      valid_cuda_versions = [101, 102, 110, 111]
      skip_on_arch(min_arch=7)
      skip_on_cuda(valid_cuda=valid_cuda_versions)
  
      Z = 2
      H = 4
      scale = 0.4
      rho = 0.4
      M = N = width
      (
          layout,
          x,
          dx,
          bool_attn_mask,
          fp_attn_mask,
          kp_mask,
      ) = init_softmax_inputs(
          Z, H, M, N, scale, rho, block, dtype, layout=None
      )
      ref_y, ref_dx = run_softmax_reference(
          x, scale, dx, kp_mask, bool_attn_mask, layout, block
      )
      st_y, st_dx = run_softmax_sparse(
          x, scale, dx, kp_mask, fp_attn_mask, layout, block
      )
  
      assert allclose(ref_y, st_y)
      assert allclose(ref_dx, st_dx)
  
file_path: DeepSpeed/tests/unit/ops/sparse_attention/test_sparse_attention.py
imports: import pytest
  import torch
  import deepspeed
  from deepspeed.accelerator import get_accelerator
  from deepspeed.ops.op_builder import SparseAttnBuilder
  from unit.util import skip_on_arch, skip_on_cuda
label_desc: Use image node_id 8 for calling a global function with example usage: test_softmax(block, width, dtype) without return types