code: def diag(v, k=0, sparse=None, gpu=None, chunk_size=None):
      """
      Extract a diagonal or construct a diagonal tensor.
  
      See the more detailed documentation for ``mt.diagonal`` if you use this
      function to extract a diagonal and wish to write to the resulting tensor
  
      Parameters
      ----------
      v : array_like
          If `v` is a 2-D tensor, return its `k`-th diagonal.
          If `v` is a 1-D tensor, return a 2-D tensor with `v` on the `k`-th
          diagonal.
      k : int, optional
          Diagonal in question. The default is 0. Use `k>0` for diagonals
          above the main diagonal, and `k<0` for diagonals below the main
          diagonal.
      sparse: bool, optional
          Create sparse tensor if True, False as default
      gpu : bool, optional
          Allocate the tensor on GPU if True, False as default
      chunk_size : int or tuple of int or tuple of ints, optional
          Desired chunk size on each dimension
  
      Returns
      -------
      out : Tensor
          The extracted diagonal or constructed diagonal tensor.
  
      See Also
      --------
      diagonal : Return specified diagonals.
      diagflat : Create a 2-D array with the flattened input as a diagonal.
      trace : Sum along diagonals.
      triu : Upper triangle of a tensor.
      tril : Lower triangle of a tensor.
  
      Examples
      --------
      >>> import mars.tensor as mt
  
      >>> x = mt.arange(9).reshape((3,3))
      >>> x.execute()
      array([[0, 1, 2],
             [3, 4, 5],
             [6, 7, 8]])
  
      >>> mt.diag(x).execute()
      array([0, 4, 8])
      >>> mt.diag(x, k=1).execute()
      array([1, 5])
      >>> mt.diag(x, k=-1).execute()
      array([3, 7])
  
      >>> mt.diag(mt.diag(x)).execute()
      array([[0, 0, 0],
             [0, 4, 0],
             [0, 0, 8]])
  
      """
      if not isinstance(v, TENSOR_TYPE):
          tensor_v = tensor(v)
          if tensor_v.issparse():
              xps = get_sparse_module(tensor_v.data)
              v = xps.csr_matrix(
                  (
                      tensor_v.op.data,
                      tensor_v.op.indices,
                      tensor_v.op.indptr,
                  ),
                  tensor_v.shape,
              )
              diag_v = sparse_diag(v, k=k)
          else:
              v = tensor(v).op.data
              diag_v = get_array_module(v).diag(v, k=k)
          sparse = sparse if sparse is not None else issparse(v)
          return tensor(
              diag_v, gpu=gpu, sparse=sparse, chunk_size=chunk_size
          )
  
      sparse = sparse if sparse is not None else v.issparse()
  
      if v.ndim == 1:
          shape = (v.size + abs(k),) * 2
      elif v.ndim == 2:
          shape = _get_diag_shape(v.shape, k)
      else:
          raise ValueError("Input must be 1- or 2-d.")
  
      op = TensorDiag(k, dtype=v.dtype, gpu=gpu, sparse=sparse)
      return op(v, shape)
  
imports: import itertools
  import numpy
  from ...None import opcodes
  from ...serialization.serializables import KeyField, Int32Field
  from ...utils import has_unknown_shape
  from ..core import TENSOR_TYPE
  from ...lib.sparse import diag
  from ...lib.sparse.core import issparse, get_array_module, get_sparse_module
  from ...lib import sparse
  from ..core import TensorOrder
  from ..array_utils import create_array
  from .core import TensorHasInput
  from .zeros import TensorZeros
  from .array import tensor
label_desc: Use image node_id 2 for calling a global function with example usage: diag(v, k, sparse, gpu, chunk_size) and returns: op, tensor