code: def test_sparse_diag_mm(create_func, sparse_shape, nnz):
      dev = F.ctx()
      diag_shape = sparse_shape[1], sparse_shape[1]
      A = create_func(sparse_shape, nnz, dev)
      diag_val = torch.randn(
          sparse_shape[1], device=dev, requires_grad=True
      )
      D = diag(diag_val, diag_shape)
      B = matmul(A, D)
      grad = torch.randn_like(B.val)
      B.val.backward(grad)
  
      torch_A = sparse_matrix_to_torch_sparse(A)
      torch_D = sparse_matrix_to_torch_sparse(D)
      torch_B = _torch_sparse_mm(torch_A, torch_D)
      torch_B_grad = sparse_matrix_to_torch_sparse(B, grad)
      torch_B.backward(torch_B_grad)
  
      with torch.no_grad():
          assert torch.allclose(
              B.to_dense(), torch_B.to_dense(), atol=1e-05
          )
          assert torch.allclose(
              val_like(A, A.val.grad).to_dense(),
              torch_A.grad.to_dense(),
              atol=1e-05,
          )
          assert torch.allclose(
              diag(D.val.grad, D.shape).to_dense(),
              torch_D.grad.to_dense(),
              atol=1e-05,
          )
  
file_path: dgl/tests/python/pytorch/sparse/test_matmul.py
imports: import warnings
  import backend
  import pytest
  import torch
  from dgl.sparse import bspmm, diag, from_coo, val_like
  from dgl.sparse.matmul import matmul
  from .utils import clone_detach_and_grad, dense_mask, rand_coo, rand_csc, rand_csr, rand_stride, sparse_matrix_to_dense, sparse_matrix_to_torch_sparse
label_desc: Use image node_id 6 for calling a global function with example usage: test_sparse_diag_mm(create_func, sparse_shape, nnz) without return types