code: def test_matmul(block, dtype, mode, trans_a, trans_b):
      valid_cuda_versions = [101, 102, 110, 111]
      skip_on_arch(min_arch=7)
      skip_on_cuda(valid_cuda=valid_cuda_versions)
  
      Z = 3
      H = 2
      M = 128
      N = 256
      K = 192
      rho = 0.5
      x, w, dy, shape, layout = init_matmul_inputs(
          Z,
          H,
          M,
          N,
          K,
          rho,
          mode,
          trans_a,
          trans_b,
          block,
          dtype,
          layout=None,
      )
      ref_y, ref_dx, ref_dw = run_matmul_reference(
          x.clone(),
          w.clone(),
          mode,
          trans_a,
          trans_b,
          layout,
          block,
          dy,
      )
      st_y, st_dx, st_dw = run_matmul_sparse(
          x.clone(),
          w.clone(),
          mode,
          trans_a,
          trans_b,
          layout,
          block,
          dy,
      )
  
      assert allclose(ref_y, st_y)
      assert allclose(ref_dx, st_dx)
      assert allclose(ref_dw, st_dw)
  
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 12 for calling a global function with example usage: test_matmul(block, dtype, mode, trans_a, trans_b) without return types