code: def diag(v_gpu):
      """
      Construct a diagonal matrix if input array is one-dimensional,
      or extracts diagonal entries of a two-dimensional array.
  
      If input-array is one-dimensional, constructs a matrix in device
      memory whose diagonal elements correspond to the elements in the
      specified array; all non-diagonal elements are set to 0.
  
      If input-array is two-dimensional, constructs an array in device memory
      whose elements correspond to the elements along the main-diagonal
      of the specified array.
  
      Parameters
      ----------
      v_obj : pycuda.gpuarray.GPUArray
              Input array of shape `(n,m)`.
  
      Returns
      -------
      d_gpu : pycuda.gpuarray.GPUArray
          If v_obj has shape `(n,1)`, output is diagonal matrix of dimensions `[n, n]`.
          If v_obj has shape `(n,m)`, output is array of length `min(n,m)`.
  
      Examples
      --------
      >>> import pycuda.driver as drv
      >>> import pycuda.gpuarray as gpuarray
      >>> import pycuda.autoinit
      >>> import numpy as np
      >>> import skcuda.linalg as linalg
      >>> linalg.init()
      >>> v = np.array([1, 2, 3, 4, 5, 6], np.float32)
      >>> v_gpu = gpuarray.to_gpu(v)
      >>> d_gpu = linalg.diag(v_gpu)
      >>> np.all(d_gpu.get() == np.diag(v))
      True
      >>> v = np.array([1j, 2j, 3j, 4j, 5j, 6j], np.complex64)
      >>> v_gpu = gpuarray.to_gpu(v)
      >>> d_gpu = linalg.diag(v_gpu)
      >>> np.all(d_gpu.get() == np.diag(v))
      True
      >>> v = np.array([[1., 2., 3.],[4., 5., 6.]], np.float64)
      >>> v_gpu = gpuarray.to_gpu(v)
      >>> d_gpu = linalg.diag(v_gpu)
      >>> d_gpu
      array([ 1.,  5.])
      """
  
      if v_gpu.dtype not in [
          np.float32,
          np.float64,
          np.complex64,
          np.complex128,
      ]:
          raise ValueError("unrecognized type")
  
      alloc = misc._global_cublas_allocator
  
      if (len(v_gpu.shape) > 1) and (len(v_gpu.shape) < 3):
          if v_gpu.dtype == np.complex64:
              func = cublas.cublasCcopy
          elif v_gpu.dtype == np.float32:
              func = cublas.cublasScopy
          elif v_gpu.dtype == np.complex128:
              func = cublas.cublasZcopy
          elif v_gpu.dtype == np.float64:
              func = cublas.cublasDcopy
          else:
              raise ValueError("unsupported input type")
  
          n = min(v_gpu.shape)
          incx = int(np.sum(v_gpu.strides) / v_gpu.dtype.itemsize)
  
          # Allocate the output array
          d_gpu = gpuarray.empty(n, v_gpu.dtype.type, allocator=alloc)
  
          handle = misc._global_cublas_handle
          func(handle, n, v_gpu.gpudata, incx, d_gpu.gpudata, 1)
          return d_gpu
      elif len(v_gpu.shape) >= 3:
          raise ValueError(
              "input array cannot have greater than 2-dimensions"
          )
  
      # Initialize output matrix:
      N = len(v_gpu)
      if N <= 0:
          raise ValueError("N must be greater than 0")
      d_gpu = misc.zeros((N, N), v_gpu.dtype, allocator=alloc)
  
      func = _get_diag_kernel(v_gpu.dtype)
      func(d_gpu, v_gpu, N, slice=slice(0, N))
      return d_gpu
  
imports: from __future__ import absolute_import, division
  from pprint import pprint
  from string import Template
  from pycuda.tools import context_dependent_memoize
  from pycuda.compiler import SourceModule
  from pycuda.reduction import ReductionKernel
  from pycuda import cumath
  import pycuda.gpuarray
  import pycuda.driver
  import pycuda.elementwise
  import pycuda.tools
  import numpy
  from .None import cublas
  from .None import cudart
  from .None import misc
  from .None import cusolver
  import sys
  from .misc import init, shutdown, add_matvec, div_matvec, mult_matvec
  from .None import install_headers
label_desc: Use image node_id 16 for calling a global function with example usage: diag(v_gpu) and returns: d_gpu, d_gpu