code: def test_diag_execution(setup):
      # 2-d  6 * 6
      a = arange(36, chunk_size=5).reshape(6, 6)
  
      d = diag(a)
      res = d.execute().fetch()
      expected = np.diag(np.arange(36).reshape(6, 6))
      np.testing.assert_equal(res, expected)
  
      d = diag(a, k=1)
      res = d.execute().fetch()
      expected = np.diag(np.arange(36).reshape(6, 6), k=1)
      np.testing.assert_equal(res, expected)
  
      d = diag(a, k=3)
      res = d.execute().fetch()
      expected = np.diag(np.arange(36).reshape(6, 6), k=3)
      np.testing.assert_equal(res, expected)
  
      d = diag(a, k=-2)
      res = d.execute().fetch()
      expected = np.diag(np.arange(36).reshape(6, 6), k=-2)
      np.testing.assert_equal(res, expected)
  
      d = diag(a, k=-5)
      res = d.execute().fetch()
      expected = np.diag(np.arange(36).reshape(6, 6), k=-5)
      np.testing.assert_equal(res, expected)
  
      # 2-d  6 * 6 sparse, no tensor
      a = sps.rand(6, 6, density=0.1)
  
      d = diag(a)
      res = d.execute().fetch()
      expected = np.diag(a.toarray())
      np.testing.assert_equal(res.toarray(), expected)
  
      d = diag(a, k=1)
      res = d.execute().fetch()
      expected = np.diag(a.toarray(), k=1)
      np.testing.assert_equal(res.toarray(), expected)
  
      d = diag(a, k=3)
      res = d.execute().fetch()
      expected = np.diag(a.toarray(), k=3)
      np.testing.assert_equal(res.toarray(), expected)
  
      d = diag(a, k=-2)
      res = d.execute().fetch()
      expected = np.diag(a.toarray(), k=-2)
      np.testing.assert_equal(res.toarray(), expected)
  
      d = diag(a, k=-5)
      res = d.execute().fetch()
      expected = np.diag(a.toarray(), k=-5)
      np.testing.assert_equal(res.toarray(), expected)
  
      # 2-d  6 * 6 sparse, from tensor
      raw_a = sps.rand(6, 6, density=0.1)
      a = tensor(raw_a, chunk_size=2)
  
      d = diag(a)
      res = d.execute().fetch()
      expected = np.diag(raw_a.toarray())
      np.testing.assert_equal(res.toarray(), expected)
  
      d = diag(a, k=1)
      res = d.execute().fetch()
      expected = np.diag(raw_a.toarray(), k=1)
      np.testing.assert_equal(res.toarray(), expected)
  
      d = diag(a, k=3)
      res = d.execute().fetch()
      expected = np.diag(raw_a.toarray(), k=3)
      np.testing.assert_equal(res.toarray(), expected)
  
      d = diag(a, k=-2)
      res = d.execute().fetch()
      expected = np.diag(raw_a.toarray(), k=-2)
      np.testing.assert_equal(res.toarray(), expected)
  
      d = diag(a, k=-5)
      res = d.execute().fetch()
      expected = np.diag(raw_a.toarray(), k=-5)
      np.testing.assert_equal(res.toarray(), expected)
  
      # 2-d  4 * 9
      a = arange(36, chunk_size=2).reshape(4, 9)
  
      d = diag(a)
      res = d.execute().fetch()
      expected = np.diag(np.arange(36).reshape(4, 9))
      np.testing.assert_equal(res, expected)
  
      d = diag(a, k=1)
      res = d.execute().fetch()
      expected = np.diag(np.arange(36).reshape(4, 9), k=1)
      np.testing.assert_equal(res, expected)
  
      d = diag(a, k=3)
      res = d.execute().fetch()
      expected = np.diag(np.arange(36).reshape(4, 9), k=3)
      np.testing.assert_equal(res, expected)
  
      d = diag(a, k=-2)
      res = d.execute().fetch()
      expected = np.diag(np.arange(36).reshape(4, 9), k=-2)
      np.testing.assert_equal(res, expected)
  
      d = diag(a, k=-3)
      res = d.execute().fetch()
      expected = np.diag(np.arange(36).reshape(4, 9), k=-3)
      np.testing.assert_equal(res, expected)
  
      # 2-d  4 * 9 sparse, no tensor
      a = sps.rand(4, 9, density=0.1)
  
      d = diag(a)
      res = d.execute().fetch()
      expected = np.diag(a.toarray())
      np.testing.assert_equal(res.toarray(), expected)
  
      d = diag(a, k=1)
      res = d.execute().fetch()
      expected = np.diag(a.toarray(), k=1)
      np.testing.assert_equal(res.toarray(), expected)
  
      d = diag(a, k=3)
      res = d.execute().fetch()
      expected = np.diag(a.toarray(), k=3)
      np.testing.assert_equal(res.toarray(), expected)
  
      d = diag(a, k=-2)
      res = d.execute().fetch()
      expected = np.diag(a.toarray(), k=-2)
      np.testing.assert_equal(res.toarray(), expected)
  
      d = diag(a, k=-3)
      res = d.execute().fetch()
      expected = np.diag(a.toarray(), k=-3)
      np.testing.assert_equal(res.toarray(), expected)
  
      # 2-d  4 * 9 sparse, from tensor
      raw_a = sps.rand(4, 9, density=0.1)
      a = tensor(raw_a, chunk_size=2)
  
      d = diag(a)
      res = d.execute().fetch()
      expected = np.diag(raw_a.toarray())
      np.testing.assert_equal(res.toarray(), expected)
  
      d = diag(a, k=1)
      res = d.execute().fetch()
      expected = np.diag(raw_a.toarray(), k=1)
      np.testing.assert_equal(res.toarray(), expected)
  
      d = diag(a, k=3)
      res = d.execute().fetch()
      expected = np.diag(raw_a.toarray(), k=3)
      np.testing.assert_equal(res.toarray(), expected)
  
      d = diag(a, k=-2)
      res = d.execute().fetch()
      expected = np.diag(raw_a.toarray(), k=-2)
      np.testing.assert_equal(res.toarray(), expected)
  
      d = diag(a, k=-3)
      res = d.execute().fetch()
      expected = np.diag(raw_a.toarray(), k=-3)
      np.testing.assert_equal(res.toarray(), expected)
  
      # 1-d
      a = arange(5, chunk_size=2)
  
      d = diag(a)
      res = d.execute().fetch()
      expected = np.diag(np.arange(5))
      np.testing.assert_equal(res, expected)
      assert res.flags["C_CONTIGUOUS"] is True
      assert res.flags["F_CONTIGUOUS"] is False
  
      d = diag(a, k=1)
      res = d.execute().fetch()
      expected = np.diag(np.arange(5), k=1)
      np.testing.assert_equal(res, expected)
  
      d = diag(a, k=3)
      res = d.execute().fetch()
      expected = np.diag(np.arange(5), k=3)
      np.testing.assert_equal(res, expected)
  
      d = diag(a, k=-2)
      res = d.execute().fetch()
      expected = np.diag(np.arange(5), k=-2)
      np.testing.assert_equal(res, expected)
  
      d = diag(a, k=-3)
      res = d.execute().fetch()
      expected = np.diag(np.arange(5), k=-3)
      np.testing.assert_equal(res, expected)
  
      d = diag(a, sparse=True)
      res = d.execute().fetch()
      expected = np.diag(np.arange(5))
      assert isinstance(res, SparseNDArray)
      np.testing.assert_equal(res.toarray(), expected)
  
      d = diag(a, k=1, sparse=True)
      res = d.execute().fetch()
      expected = np.diag(np.arange(5), k=1)
      assert isinstance(res, SparseNDArray)
      np.testing.assert_equal(res.toarray(), expected)
  
      d = diag(a, k=2, sparse=True)
      res = d.execute().fetch()
      expected = np.diag(np.arange(5), k=2)
      assert isinstance(res, SparseNDArray)
      np.testing.assert_equal(res.toarray(), expected)
  
      d = diag(a, k=-2, sparse=True)
      res = d.execute().fetch()
      expected = np.diag(np.arange(5), k=-2)
      assert isinstance(res, SparseNDArray)
      np.testing.assert_equal(res.toarray(), expected)
  
      d = diag(a, k=-3, sparse=True)
      res = d.execute().fetch()
      expected = np.diag(np.arange(5), k=-3)
      assert isinstance(res, SparseNDArray)
      np.testing.assert_equal(res.toarray(), expected)
  
imports: import tempfile
  import shutil
  import os
  import time
  import numpy
  import pandas
  import pytest
  from ....None import tensor
  from ....None import dataframe
  from ....lib.sparse import SparseNDArray
  from ....tests.core import require_cupy
  from ....utils import lazy_import
  from ...lib import nd_grid
  from ..None import tensor, ones_like, zeros, zeros_like, full, full_like, arange, empty, empty_like, diag, diagflat, eye, linspace, meshgrid, indices, triu, tril, from_dataframe, fromtiledb, fromhdf5, fromzarr
label_desc: Use image node_id 7 for calling a global function with example usage: test_diag_execution(setup) without return types