code: def init_softmax_inputs(
      Z, H, M, N, scale, rho, block, dtype, dense_x=True, layout=None
  ):
      if layout is None:
          layout = make_layout(rho, (H, M // block, N // block))
      if dense_x:
          x = torch.rand(
              (Z, H, M, N),
              dtype=dtype,
              requires_grad=True,
              device=get_accelerator().device_name(),
          )
      else:
          x = torch.rand(
              (Z, layout.sum(), block, block),
              dtype=dtype,
              requires_grad=True,
              device=get_accelerator().device_name(),
          )
      dx = torch.rand_like(x)
      bool_attn_mask = torch.randint(
          low=0,
          high=2,
          size=(N, N),
          dtype=torch.bool,
          requires_grad=False,
          device=get_accelerator().device_name(),
      )
      fp_attn_mask = bool_attn_mask.type(dtype)
      kp_mask = torch.randint(
          low=0,
          high=2,
          size=(Z, N),
          dtype=dtype,
          requires_grad=False,
          device=get_accelerator().device_name(),
      )
      kp_mask[kp_mask == 1.0] = float("-inf")
      return layout, x, dx, bool_attn_mask, fp_attn_mask, kp_mask
  
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 7 for calling a global function with example usage: init_softmax_inputs(Z, H, M, N, scale, rho, block, dtype, dense_x, layout) and returns: layout,  x,  dx,  bool_attn_mask,  fp_attn_mask,  kp_mask