code: def test_matmul_with_mask(device, contiguous, is_sparse):
      B, L, K = 8, 30, 32
      prob = 0.5
      a = torch.rand(B, L, K, device=device)
      b = torch.rand(B, K, L, device=device)
      if not contiguous:
          a = a.transpose(-2, -1).contiguous().transpose(-2, -1)
          b = b.transpose(-2, -1).contiguous().transpose(-2, -1)
      mask = torch.rand(B, L, L, device=device) > prob
  
      fn = torch.ops.xformers.matmul_with_mask
      fn_gt = _baseline_matmul_with_dense_mask
  
      if is_sparse:
          mask = mask.to_sparse()
          fn_gt = _baseline_matmul_with_sparse_mask
  
      res = fn(a, b, mask)
      res_gt = fn_gt(a, b, mask)
  
      if is_sparse:
          res = res.to_dense()
          res_gt = res_gt.to_dense()
  
      assert res.dtype == res_gt.dtype
      assert torch.allclose(res, res_gt)
  
imports: import pytest
  import torch
  import xformers
  import xformers.components.attention.core
  from xformers.components.attention._sputnik_sparse import _csr_to_coo
  from xformers.components.attention.core import _broadcast_batch, _create_random_sparsity, _sparse_bmm
label_desc: Use image node_id 4 for calling a global function with example usage: test_matmul_with_mask(device, contiguous, is_sparse) without return types