def METHOD_NAME(create_func, sparse_shape, nnz):
    dev = F.ctx()
    diag_shape = sparse_shape[0], sparse_shape[0]
    A = create_func(sparse_shape, nnz, dev)
    diag_val = torch.randn(sparse_shape[0], device=dev, requires_grad=True)
    D = diag(diag_val, diag_shape)
    B = matmul(D, A)
    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_D, torch_A)
    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,
        )