code: def test_diag(val_shape, mat_shape):
      ctx = F.ctx()
      # creation
      val = torch.randn(val_shape).to(ctx)
      mat = diag(val, mat_shape)
  
      # val, shape attributes
      assert torch.allclose(mat.val, val)
      if mat_shape is None:
          mat_shape = (val_shape[0], val_shape[0])
      assert mat.shape == mat_shape
  
      val = torch.randn(val_shape).to(ctx)
  
      # nnz
      assert mat.nnz == val.shape[0]
      # dtype
      assert mat.dtype == val.dtype
      # device
      assert mat.device == val.device
  
      # row, col, val
      edge_index = torch.arange(len(val)).to(mat.device)
      row, col = mat.coo()
      val = mat.val
      assert torch.allclose(row, edge_index)
      assert torch.allclose(col, edge_index)
      assert torch.allclose(val, val)
  
imports: import unittest
  import warnings
  import backend
  import pytest
  import torch
  from dgl.sparse import diag, from_coo, from_csc, from_csr, from_torch_sparse, identity, to_torch_sparse_coo, to_torch_sparse_csc, to_torch_sparse_csr, val_like
  from .utils import rand_coo, rand_csc, rand_csr, rand_diag, sparse_matrix_to_dense
label_desc: Use image node_id 28 for calling a global function with example usage: test_diag(val_shape, mat_shape) without return types