code: def diag(v, k=0):
      if not isinstance(v, np.ndarray) and not isinstance(v, Array):
          raise TypeError(
              f"v must be a dask array or numpy array, got {type(v)}"
          )
  
      name = "diag-" + tokenize(v, k)
  
      meta = meta_from_array(v, 2 if v.ndim == 1 else 1)
  
      if isinstance(v, np.ndarray) or (
          hasattr(v, "__array_function__") and not isinstance(v, Array)
      ):
          if v.ndim == 1:
              m = abs(k)
              chunks = ((v.shape[0] + m,), (v.shape[0] + m,))
              dsk = {(name, 0, 0): (np.diag, v, k)}
          elif v.ndim == 2:
              kdiag_row_start = max(0, -k)
              kdiag_row_stop = min(v.shape[0], v.shape[1] - k)
              len_kdiag = kdiag_row_stop - kdiag_row_start
              chunks = ((0,),) if len_kdiag <= 0 else ((len_kdiag,),)
              dsk = {(name, 0): (np.diag, v, k)}
          else:
              raise ValueError("Array must be 1d or 2d only")
          return Array(dsk, name, chunks, meta=meta)
  
      if v.ndim != 1:
          if v.ndim != 2:
              raise ValueError("Array must be 1d or 2d only")
          if k == 0 and v.chunks[0] == v.chunks[1]:
              dsk = {
                  (name, i): (np.diag, row[i])
                  for i, row in enumerate(v.__dask_keys__())
              }
              graph = HighLevelGraph.from_collections(
                  name, dsk, dependencies=[v]
              )
              return Array(graph, name, (v.chunks[0],), meta=meta)
          else:
              return diagonal(v, k)
  
      if k == 0:
          chunks_1d = v.chunks[0]
          blocks = v.__dask_keys__()
          dsk = {}
          for i, m in enumerate(chunks_1d):
              for j, n in enumerate(chunks_1d):
                  key = (name, i, j)
                  if i == j:
                      dsk[key] = (np.diag, blocks[i])
                  else:
                      dsk[key] = (np.zeros, (m, n))
                      dsk[key] = (
                          partial(np.zeros_like, shape=(m, n)),
                          meta,
                      )
  
          graph = HighLevelGraph.from_collections(
              name, dsk, dependencies=[v]
          )
          return Array(graph, name, (chunks_1d, chunks_1d), meta=meta)
  
      elif k > 0:
          return pad(diag(v), [[0, k], [k, 0]], mode="constant")
      elif k < 0:
          return pad(diag(v), [[-k, 0], [0, -k]], mode="constant")
  
imports: from __future__ import annotations
  import itertools
  from collections.abc import Sequence
  from functools import partial
  from itertools import product
  from numbers import Integral, Number
  import numpy
  from tlz import sliding_window
  from dask.array import chunk
  from dask.array.backends import array_creation_dispatch
  from dask.array.core import Array, asarray, block, blockwise, broadcast_arrays, broadcast_to, concatenate, normalize_chunks, stack
  from dask.array.numpy_compat import AxisError
  from dask.array.ufunc import greater_equal, rint
  from dask.array.utils import meta_from_array
  from dask.array.wrap import empty, full, ones, zeros
  from dask.base import tokenize
  from dask.highlevelgraph import HighLevelGraph
  from dask.utils import cached_cumsum, derived_from, is_cupy_type
label_desc: Use image node_id 12 for calling a global function with example usage: diag(v, k) and returns: Array, Array, Array, diagonal, pad, pad