code: def diag_sparse(A):
      
      if isspmatrix(A):
          return A.diagonal()
      else:
          if(np.ndim(A) != 1):
              raise ValueError('input diagonal array expected to be 1d')
          return csr_matrix((np.asarray(A), np.arange(len(A)),
                             np.arange(len(A)+1)), (len(A), len(A)))
docs: Return a diagonal.
  
      If A is a sparse matrix (e.g. csr_matrix or csc_matrix)
         - return the diagonal of A as an array
  
      Otherwise
         - return a csr_matrix with A on the diagonal
  
      Parameters
      ----------
      A : sparse matrix or 1d array
          General sparse matrix or array of diagonal entries
  
      Returns
      -------
      B : array or sparse matrix
          Diagonal sparse is returned as csr if A is dense otherwise return an
          array of the diagonal
  
      Examples
      --------
      >>> import numpy as np
      >>> from pyamg.util.utils import diag_sparse
      >>> d = 2.0*np.ones((3,)).ravel()
      >>> print diag_sparse(d).todense()
      [[ 2.  0.  0.]
       [ 0.  2.  0.]
       [ 0.  0.  2.]]