code: def test_masked_matmul(tensor_type, device):
      N, C, H, W, L = 8, 2, 64, 64, 32
      sparsity = 0.7
      dtype, atol = _get_dtype_atol(tensor_type, device)
  
      shape0 = (N, C, H, W)
      shape1 = (N, C, H, L)
      shape2 = (N, C, W, L)
  
      if tensor_type != BlockSparseTensor:
          shape0 = shape0[1:]
          shape1 = shape1[1:]
          shape2 = shape2[1:]
  
      mask_sparse = _create_tensor(
          tensor_type,
          device,
          dtype=torch.bool,
          shape=shape0,
          sparsity=sparsity,
      )
      mask = mask_sparse.to_dense()
  
      a = torch.randn(shape1, device=device, dtype=dtype)
      b = torch.randn(shape2, device=device, dtype=dtype)
  
      aa = a.clone()
      bb = b.clone()
  
      a.requires_grad_(True)
      b.requires_grad_(True)
      aa.requires_grad_(True)
      bb.requires_grad_(True)
  
      bt = b.transpose(-2, -1)
      bbt = bb.transpose(-2, -1)
  
      res_gt = masked_matmul(a, bt, mask)
      res = masked_matmul(aa, bbt, mask_sparse)
  
      res_dense = res.to_dense()
      res_dense = torch.where(
          mask, res_dense, torch.full_like(res_dense, float("-inf"))
      )
  
      assert res.dtype == res_gt.dtype
      assert torch.allclose(res_dense, res_gt, atol=atol)
  
      # try to workaround non-contiguous issues with triton for now
      res_gt.backward(torch.ones_like(res_gt))
      res.values().backward(torch.ones_like(res.values()))
  
      assert torch.allclose(a.grad, aa.grad, atol=atol)
      assert torch.allclose(b.grad, bb.grad, atol=atol)
  
file_path: xformers/tests/test_sparse_tensors.py
imports: import pytest
  import torch
  import xformers
  from xformers.ops import masked_matmul
  from xformers.sparse import BlockSparseTensor, SparseCSRTensor
label_desc: Use image node_id 7 for calling a global function with example usage: test_masked_matmul(tensor_type, device) without return types