code: def test_diag_extraction(k):
      # when input 2d-array is a square numpy array:
      x = np.arange(64).reshape((8, 8))
      assert_eq(da.diag(x, k), np.diag(x, k))
      # when input 2d-array is a square dask array:
      d = da.from_array(x, chunks=(4, 4))
      assert_eq(da.diag(d, k), np.diag(x, k))
      # heterogeneous chunks:
      d = da.from_array(x, chunks=((3, 2, 3), (4, 1, 2, 1)))
      assert_eq(da.diag(d, k), np.diag(x, k))
  
      # when input 2d-array is a rectangular numpy array:
      y = np.arange(5 * 8).reshape((5, 8))
      assert_eq(da.diag(y, k), np.diag(y, k))
      # when input 2d-array is a rectangular dask array:
      d = da.from_array(y, chunks=(4, 4))
      assert_eq(da.diag(d, k), np.diag(y, k))
      # heterogeneous chunks:
      d = da.from_array(y, chunks=((3, 2), (4, 1, 2, 1)))
      assert_eq(da.diag(d, k), np.diag(y, k))
  
imports: from __future__ import annotations
  import pytest
  import numpy
  import pytest
  from tlz import concat
  import dask
  import dask.array
  from dask.array.core import normalize_chunks
  from dask.array.numpy_compat import AxisError
  from dask.array.utils import assert_eq, same_keys
label_desc: Use image node_id 18 for calling a global function with example usage: test_diag_extraction(k) without return types