code: def test_diag():
      # test 2-d, shape[0] == shape[1], k == 0
      v = tensor(np.arange(16).reshape(4, 4), chunk_size=2)
      t = diag(v)
  
      assert t.shape == (4,)
      assert t.op.gpu is None
      t = tile(t)
      assert t.nsplits == ((2, 2),)
  
      v = tensor(np.arange(16).reshape(4, 4), chunk_size=(2, 3))
      t = diag(v)
  
      assert t.shape == (4,)
      t = tile(t)
      assert t.nsplits == ((2, 1, 1),)
  
      # test 1-d, k == 0
      v = tensor(np.arange(3), chunk_size=2)
      t = diag(v, sparse=True)
  
      assert t.shape == (3, 3)
      t = tile(t)
      assert t.nsplits == ((2, 1), (2, 1))
      assert (
          len(
              [
                  c
                  for c in t.chunks
                  if c.op.__class__.__name__ == "TensorDiag"
              ]
          )
          == 2
      )
      assert t.chunks[0].op.sparse is True
  
      # test 2-d, shape[0] != shape[1]
      v = tensor(np.arange(24).reshape(4, 6), chunk_size=2)
      t = diag(v)
  
      assert t.shape == np.diag(np.arange(24).reshape(4, 6)).shape
      t = tile(t)
      assert tuple(sum(s) for s in t.nsplits) == t.shape
  
      v = tensor(np.arange(24).reshape(4, 6), chunk_size=2)
  
      t = diag(v, k=1)
      assert t.shape == np.diag(np.arange(24).reshape(4, 6), k=1).shape
      t = tile(t)
      assert tuple(sum(s) for s in t.nsplits) == t.shape
  
      t = diag(v, k=2)
      assert t.shape == np.diag(np.arange(24).reshape(4, 6), k=2).shape
      t = tile(t)
      assert tuple(sum(s) for s in t.nsplits) == t.shape
  
      t = diag(v, k=-1)
      assert t.shape == np.diag(np.arange(24).reshape(4, 6), k=-1).shape
      t = tile(t)
      assert tuple(sum(s) for s in t.nsplits) == t.shape
  
      t = diag(v, k=-2)
      assert t.shape == np.diag(np.arange(24).reshape(4, 6), k=-2).shape
      t = tile(t)
      assert tuple(sum(s) for s in t.nsplits) == t.shape
  
      # test tiled zeros' keys
      a = arange(5, chunk_size=2)
      t = diag(a)
      t = tile(t)
      # 1 and 2 of t.chunks is ones, they have different shapes
      assert t.chunks[1].op.key != t.chunks[2].op.key
  
imports: import shutil
  import tempfile
  from copy import copy
  import numpy
  import scipy.sparse
  import pytest
  from ....None import dataframe
  from ....core import enter_mode, tile
  from ...None import ones, zeros, tensor, full, arange, diag, linspace, triu, tril, ones_like
  from ...core import Tensor, SparseTensor
  from ..None import array, fromtiledb, TensorTileDBDataSource, fromdense, asarray, ascontiguousarray, asfortranarray
  from ..tri import TensorTriu, TensorTril
  from ..zeros import TensorZeros
  from ..from_dense import DenseToSparse
  from ..array import CSRMatrixDataSource
  from ..ones import TensorOnes, TensorOnesLike
  from ..from_dataframe import from_dataframe
label_desc: Use image node_id 9 for calling a global function with example usage: test_diag() without return types