code: def run_softmax_reference(
      x, scale, dx, kp_mask, attn_mask, layout, block
  ):
      x = sparse_to_dense(x, layout, block, zero=float("-inf"))
      x.retain_grad()
      if kp_mask is not None:
          bcattn_mask = attn_mask[None, None, :, :] + torch.zeros_like(
              x
          )
          x[bcattn_mask == 0] = float("-inf")
          y = torch.softmax(x * scale + kp_mask[:, None, None, :], -1)
      else:
          y = torch.softmax(x * scale, -1)
      y.backward(dx)
      dx = x.grad.clone()
      dx = dense_to_sparse(dx, layout, block)
      y = dense_to_sparse(y, layout, block)
      return y, 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 5 for calling a global function with example usage: run_softmax_reference(x, scale, dx, kp_mask, attn_mask, layout, block) and returns: y,  dx