code: def diag_sparse_matrix(v, k=0, gpu=False):
      v = naked(v)
      if gpu and get_array_module(v) is not cp:
          v = cp.asarray(v)
      if not gpu and get_array_module(v) is not np:
          v = v.get()
  
      if v.ndim == 1:
          sparse_m = sps if not gpu else cps
          m = n = v.size + k
          mat = sparse_m.spdiags(v[None], [k], m, n, format="csr")
          return SparseMatrix(mat)
      else:
          assert v.ndim == 2
          sparse_m = sps if not gpu else cps
          sparse_eye = sparse_m.eye(v.shape[0], v.shape[1], k=k)
          mat = sparse_eye.multiply(v).tocoo()
          size = sparse_eye.nnz
          col = mat.col - max(k, 0)
          row = get_array_module(col).zeros((len(col),))
          return SparseNDArray(
              sparse_m.csr_matrix(
                  (mat.data, (row, col)), shape=(1, size)
              ),
              shape=(size,),
          )
  
imports: import numpy
  from collections.abc import Iterable
  from typing import List
  from .core import issparse, get_array_module, cp, cps, get_sparse_module, naked, sps, splinalg
  from .array import SparseNDArray, SparseArray
label_desc: Use image node_id 2 for calling a global function with example usage: diag_sparse_matrix(v, k, gpu) and returns: SparseMatrix, SparseNDArray