<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 def diag_sparse(A):
 """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 -------- [[ 2. 0. 0.] [ 0. 2. 0.] [ 0. 0. 2.]] """

    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)))