code: def test_sddmm_coo(L, M, K, prob):
      device = torch.device("cuda")
      # TODO add more checks for different nnz
      B = 8
      a = torch.rand(B, L, K, device=device)
      b = torch.rand(B, M, K, device=device)
      mask = _create_random_sparsity(
          torch.ones(B, L, M, dtype=torch.bool, device=device), prob
      )
  
      mask_csr = xformers.components.attention.core.SparseCS(
          mask, device
      )
      row_indices = mask_csr.row_indices
      row_offsets = mask_csr.row_offsets
      column_indices = mask_csr.column_indices
  
      fn = torch.ops.xformers.coo_sddmm
      fn_gt = torch.ops.xformers.sddmm_sputnik
  
      # convert from csr to coo
      row_coo, _ = _csr_to_coo(L, M, row_offsets, column_indices)
  
      res = fn(a, b, row_indices, row_coo, column_indices)
      res_gt = fn_gt(a, b, row_indices, row_offsets, column_indices)
  
      assert res.dtype == res_gt.dtype
      assert torch.allclose(res, res_gt, atol=1e-6)
  
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 9 for calling a global function with example usage: test_sddmm_coo(L, M, K, prob) without return types