code: def test_sparse_softmax(tensor_type, device):
      N, C, H, W = 8, 2, 64, 64
      dtype, atol = _get_dtype_atol(tensor_type, device)
  
      sparsity = 0.8
  
      shape0 = (N, C, H, W)
      if tensor_type != BlockSparseTensor:
          shape0 = shape0[1:]
  
      a_sparse = _create_tensor(
          tensor_type,
          device,
          dtype=dtype,
          shape=shape0,
          sparsity=sparsity,
      )
      a = a_sparse.to_dense()
      mask = a != 0
  
      a[~mask] = float("-inf")
  
      a_sparse.requires_grad_(True)
      a.requires_grad_(True)
  
      res_gt = torch.softmax(a, dim=-1)
      res_sparse = torch.softmax(a_sparse, dim=-1)
  
      res = res_sparse.to_dense()
  
      assert res.dtype == res_gt.dtype
      assert torch.allclose(
          res, res_gt, atol=atol
      ), f"{torch.max(torch.abs(res- res_gt))}"
  
      # WARNING: gradients are modified in-place!
      res_sparse.values().backward(torch.ones_like(res_sparse.values()))
      res_gt.backward(torch.ones_like(res_gt))
  
      a_grad = a.grad.clone()
      a_grad[~mask] = 0
  
      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 9 for calling a global function with example usage: test_sparse_softmax(tensor_type, device) without return types