(v_gpu)
return d_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.])
Construct a diagonal matrix if input array is one-dimensional,
    or extracts diagonal entries of a two-dimensional array.
Construct a diagonal matrix if input array is one - dimensional or extracts diagonal entries of a two - dimensional array .
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
def diag ( v_gpu ) : 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