code: def test_bmm(tensor_type, device):
      N, C, H, W, L = 8, 2, 64, 64, 32
      dtype, atol = _get_dtype_atol(tensor_type, device)
  
      sparsity = 0.8
      shape0 = (N, C, H, W)
      shape1 = (N, C, W, L)
  
      if tensor_type != BlockSparseTensor:
          shape0 = shape0[1:]
          shape1 = shape1[1:]
  
      a_sparse = _create_tensor(
          tensor_type,
          device,
          dtype=dtype,
          shape=shape0,
          sparsity=sparsity,
      )
      a = a_sparse.to_dense()
      mask = a != 0
  
      a_sparse.requires_grad_(True)
      a.requires_grad_(True)
  
      b = torch.randn(shape1, device=device, dtype=dtype)
      b2 = b.clone()
  
      b.requires_grad_(True)
      b2.requires_grad_(True)
  
      res_gt = a @ b
      res = a_sparse @ b2
  
      assert res.dtype == res_gt.dtype
      assert torch.allclose(
          res, res_gt, atol=atol
      ), f"{torch.max(torch.abs(res-res_gt))} - tolerance: {atol}"
  
      res_gt.sum().backward()
      res.sum().backward()
  
      a_grad = a.grad.clone().detach()
      a_grad[~mask] = 0
  
      assert torch.allclose(b.grad, b2.grad, atol=atol)
      assert torch.allclose(
          a_grad, a_sparse.grad.to_dense(), atol=atol
      ), f"{torch.max(torch.abs(a_grad-a_sparse.grad.to_dense()))}"
  
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 8 for calling a global function with example usage: test_bmm(tensor_type, device) without return types