code: def truncate_rows(A, nz_per_row):
      
      if not isspmatrix(A):
          raise ValueError("Sparse matrix input needed")
      if isspmatrix_bsr(A):
          blocksize = A.blocksize
      if isspmatrix_csr(A):
          A = A.copy()    # don't modify A in-place
      Aformat = A.format
      A = A.tocsr()
      nz_per_row = int(nz_per_row)
  
      # Truncate rows of A, and then convert A back to original format
      pyamg.amg_core.truncate_rows_csr(A.shape[0], nz_per_row, A.indptr,
                                       A.indices, A.data)
  
      A.eliminate_zeros()
      if Aformat == 'bsr':
          A = A.tobsr(blocksize)
      else:
          A = A.asformat(Aformat)
  
      return A
docs: Truncate the rows of A by keeping only the largest in magnitude entries in each row.
  
      Parameters
      ----------
      A : sparse_matrix
  
      nz_per_row : int
          Determines how many entries in each row to keep
  
      Returns
      -------
      A : sparse_matrix
          Each row has been truncated to at most nz_per_row entries
  
      Examples
      --------
      >>> from pyamg.gallery import poisson
      >>> from pyamg.util.utils import truncate_rows
      >>> from scipy import array
      >>> from scipy.sparse import csr_matrix
      >>> A = csr_matrix( array([[-0.24, -0.5 ,  0.  ,  0.  ],
      ...                        [ 1.  , -1.1 ,  0.49,  0.1 ],
      ...                        [ 0.  ,  0.4 ,  1.  ,  0.5 ]])  )
      >>> truncate_rows(A, 2).todense()
      matrix([[-0.24, -0.5 ,  0.  ,  0.  ],
              [ 1.  , -1.1 ,  0.  ,  0.  ],
              [ 0.  ,  0.  ,  1.  ,  0.5 ]])