Repository: cai4cai/torchsparsegradutils (language: python)

Issue:
Provide LSMR with no extra dependency
This is a follow-up from #3 

Options include adapting code of:
- LSMR from DINGO https://github.com/RixonC/DINGO
- LSMR from pytorch-minimize: https://github.com/rfeinman/pytorch-minimize/blob/master/torchmin/lstsq/lsmr.py


Discussion / hints:

Fix patch:
diff --git a/torchsparsegradutils/utils/__init__.py b/torchsparsegradutils/utils/__init__.py
index e798bea..1bcb4f7 100644
--- a/torchsparsegradutils/utils/__init__.py
+++ b/torchsparsegradutils/utils/__init__.py
@@ -1,5 +1,6 @@
 from .linear_cg import linear_cg, LinearCGSettings
 from .minres import minres, MINRESSettings
 from .bicgstab import bicgstab, BICGSTABSettings
+from .lsmr import lsmr
 
-__all__ = ["linear_cg", "LinearCGSettings", "minres", "MINRESSettings", "bicgstab", "BICGSTABSettings"]
+__all__ = ["linear_cg", "LinearCGSettings", "minres", "MINRESSettings", "bicgstab", "BICGSTABSettings", "lsmr"]
diff --git a/torchsparsegradutils/utils/bicgstab.py b/torchsparsegradutils/utils/bicgstab.py
index 5013730..ec14dbb 100644
--- a/torchsparsegradutils/utils/bicgstab.py
+++ b/torchsparsegradutils/utils/bicgstab.py
@@ -88,9 +88,9 @@ def bicgstab(
 
     if initial_guess is None:
         x = torch.zeros(n, dtype=res_dtype, device=res_device)
-    else
+    else:
         x = initial_guess.clone()
-    
+
     # matvec_max = kwargs.get('matvec_max', 2*n)
     matvec_max = 2 * n if settings.matvec_max is None else settings.matvec_max
 
diff --git a/torchsparsegradutils/utils/lsmr.py b/torchsparsegradutils/utils/lsmr.py
new file mode 100644
index 0000000..b89c17e
--- /dev/null
+++ b/torchsparsegradutils/utils/lsmr.py
@@ -0,0 +1,324 @@
+"""
+Code adapted from https://github.com/rfeinman/pytorch-minimize/blob/master/torchmin/lstsq/lsmr.py
+Code modified from scipy.sparse.linalg.lsmr
+
+Copyright (C) 2010 David Fong and Michael Saunders
+"""
+import torch
+
+
+def _sym_ortho(a, b, out):
+    torch.hypot(a, b, out=out[2])
+    torch.div(a, out[2], out=out[0])
+    torch.div(b, out[2], out=out[1])
+    return out
+
+
+@torch.no_grad()
+def lsmr(
+    A, b, Armat=None, n=None, damp=0.0, atol=1e-6, btol=1e-6, conlim=1e8, maxiter=None, x0=None, check_nonzero=True
+):
+    """Iterative solver for least-squares problems.
+
+    lsmr solves the system of linear equations ``Ax = b``. If the system
+    is inconsistent, it solves the least-squares problem ``min ||b - Ax||_2``.
+    ``A`` is a rectangular matrix of dimension m-by-n, where all cases are
+    allowed: m = n, m > n, or m < n. ``b`` is a vector of length m.
+    The matrix A may be dense or sparse (usually sparse).
+
+    Parameters
+    ----------
+    A : {matrix, sparse matrix, ndarray, LinearOperator}
+        Matrix A in the linear system.
+        Alternatively, ``A`` can be a linear operator which can
+        produce ``Ax`` and ``A^H x`` using, e.g.,
+        ``scipy.sparse.linalg.LinearOperator``.
+    b : array_like, shape (m,)
+        Vector ``b`` in the linear system.
+    damp : float
+        Damping factor for regularized least-squares. `lsmr` solves
+        the regularized least-squares problem::
+         min ||(b) - (  A   )x||
+             ||(0)   (damp*I) ||_2
+        where damp is a scalar.  If damp is None or 0, the system
+        is solved without regularization.
+    atol, btol : float, optional
+        Stopping tolerances. `lsmr` continues iterations until a
+        certain backward error estimate is smaller than some quantity
+        depending on atol and btol.  Let ``r = b - Ax`` be the
+        residual vector for the current approximate solution ``x``.
+        If ``Ax = b`` seems to be consistent, ``lsmr`` terminates
+        when ``norm(r) <= atol * norm(A) * norm(x) + btol * norm(b)``.
+        Otherwise, lsmr terminates when ``norm(A^H r) <=
+        atol * norm(A) * norm(r)``.  If both tolerances are 1.0e-6 (say),
+        the final ``norm(r)`` should be accurate to about 6
+        digits. (The final ``x`` will usually have fewer correct digits,
+        depending on ``cond(A)`` and the size of LAMBDA.)  If `atol`
+        or `btol` is None, a default value of 1.0e-6 will be used.
+        Ideally, they should be estimates of the relative error in the
+        entries of ``A`` and ``b`` respectively.  For example, if the entries
+        of ``A`` have 7 correct digits, set ``atol = 1e-7``. This prevents
+        the algorithm from doing unnecessary work beyond the
+        uncertainty of the input data.
+    conlim : float, optional
+        `lsmr` terminates if an estimate of ``cond(A)`` exceeds
+        `conlim`.  For compatible systems ``Ax = b``, conlim could be
+        as large as 1.0e+12 (say).  For least-squares problems,
+        `conlim` should be less than 1.0e+8. If `conlim` is None, the
+        default value is 1e+8.  Maximum precision can be obtained by
+        setting ``atol = btol = conlim = 0``, but the number of
+        iterations may then be excessive.
+    maxiter : int, optional
+        `lsmr` terminates if the number of iterations reaches
+        `maxiter`.  The default is ``maxiter = min(m, n)``.  For
+        ill-conditioned systems, a larger value of `maxiter` may be
+        needed.
+    x0 : array_like, shape (n,), optional
+        Initial guess of ``x``, if None zeros are used.
+
+    Returns
+    -------
+    x : ndarray of float
+        Least-square solution returned.
+    itn : int
+        Number of iterations used.
+
+    """
+    if torch.is_tensor(A):
+        if n is None:
+            n = A.shape[1]
+        if Armat is None:
+            Armat = (torch.adjoint(A)).matmul
+        A = A.matmul
+    elif not callable(A):
+        raise RuntimeError("matmul_closure must be a tensor, or a callable object!")
+
+    if n is None:
+        raise RuntimeError("n needs to be provided of computed from A given as a tensor")
+
+    if torch.is_tensor(Armat):
+        Armat = Armat.matmul
+    elif not callable(Armat):
+        raise RuntimeError("matmul_closure must be a tensor, or a callable object!")
+
+    sdtype = b.dtype
+    if b.dtype == torch.complex64:
+        sdtype = torch.float32
+    elif b.dtype == torch.complex128:
+        sdtype = torch.float64
+
+    b = torch.atleast_1d(b)
+    if b.dim() > 1:
+        b = b.squeeze()
+    eps = torch.finfo(sdtype).eps
+    damp = torch.as_tensor(damp, dtype=sdtype, device=b.device)
+    ctol = 1 / conlim if conlim > 0 else 0.0
+    m = b.shape[0]
+    if maxiter is None:
+        maxiter = min(m, n)
+
+    u = b.clone()
+    normb = b.norm()
+    if x0 is None:
+        x = b.new_zeros(n)
+        beta = normb.clone()
+    else:
+        x = torch.atleast_1d(x0).clone()
+        u.sub_(A(x))
+        beta = u.norm()
+
+    if beta > 0:
+        u.div_(beta)
+        v = Armat(u)
+        alpha = v.norm()
+    else:
+        v = b.new_zeros(n)
+        alpha = b.new_tensor(0, dtype=sdtype)
+
+    v = torch.where(alpha > 0, v / alpha, v)
+
+    # Initialize variables for 1st iteration.
+
+    zetabar = alpha * beta
+    alphabar = alpha.clone()
+    rho = b.new_tensor(1, dtype=sdtype)
+    rhobar = b.new_tensor(1, dtype=sdtype)
+    cbar = b.new_tensor(1, dtype=sdtype)
+    sbar = b.new_tensor(0, dtype=sdtype)
+
+    h = v.clone()
+    hbar = b.new_zeros(n)
+
+    # Initialize variables for estimation of ||r||.
+
+    betadd = beta.clone()
+    betad = b.new_tensor(0, dtype=sdtype)
+    rhodold = b.new_tensor(1, dtype=sdtype)
+    tautildeold = b.new_tensor(0, dtype=sdtype)
+    thetatilde = b.new_tensor(0, dtype=sdtype)
+    zeta = b.new_tensor(0, dtype=sdtype)
+    d = b.new_tensor(0, dtype=sdtype)
+
+    # Initialize variables for estimation of ||A|| and cond(A)
+
+    normA2 = alpha.square()
+    maxrbar = b.new_tensor(0, dtype=sdtype)
+    minrbar = b.new_tensor(0.99 * torch.finfo(sdtype).max, dtype=sdtype)
+    normA = normA2.sqrt()
+    condA = b.new_tensor(1, dtype=sdtype)
+    normx = b.new_tensor(0, dtype=sdtype)
+    # normar = b.new_tensor(0,dtype=sdtype)
+    # normr = b.new_tensor(0,dtype=sdtype)
+
+    normr = beta.clone()
+    normar = alpha * beta
+    if normar == 0:
+        return x, 0
+
+    if normb == 0:
+        x[:] = 0
+        return x, 0
+
+    # extra buffers (added by Reuben)
+    c = b.new_tensor(0, dtype=sdtype)
+    s = b.new_tensor(0, dtype=sdtype)
+    chat = b.new_tensor(0, dtype=sdtype)
+    shat = b.new_tensor(0, dtype=sdtype)
+    alphahat = b.new_tensor(0, dtype=sdtype)
+    ctildeold = b.new_tensor(0, dtype=sdtype)
+    stildeold = b.new_tensor(0, dtype=sdtype)
+    rhotildeold = b.new_tensor(0, dtype=sdtype)
+    rhoold = b.new_tensor(0, dtype=sdtype)
+    rhobarold = b.new_tensor(0, dtype=sdtype)
+    zetaold = b.new_tensor(0, dtype=sdtype)
+    thetatildeold = b.new_tensor(0, dtype=sdtype)
+    betaacute = b.new_tensor(0, dtype=sdtype)
+    betahat = b.new_tensor(0, dtype=sdtype)
+    betacheck = b.new_tensor(0, dtype=sdtype)
+    taud = b.new_tensor(0, dtype=sdtype)
+
+    # Main iteration loop.
+    for itn in range(1, maxiter + 1):
+
+        # Perform the next step of the bidiagonalization to obtain the
+        # next  beta, u, alpha, v.  These satisfy the relations
+        #         beta*u  =  a*v   -  alpha*u,
+        #        alpha*v  =  A'*u  -  beta*v.
+
+        u.mul_(-alpha).add_(A(v))
+        torch.norm(u, out=beta)
+
+        if (not check_nonzero) or beta > 0:
+            # check_nonzero option provides a means to avoid the GPU-CPU
+            # synchronization of a `beta > 0` check. For most cases
+            # beta == 0 is unlikely, but use this option with caution.
+            u.div_(beta)
+            v.mul_(-beta).add_(Armat(u))
+            torch.norm(v, out=alpha)
+            v = torch.where(alpha > 0, v / alpha, v)
+
+        # At this point, beta = beta_{k+1}, alpha = alpha_{k+1}.
+
+        _sym_ortho(alphabar, damp, out=(chat, shat, alphahat))
+
+        # Use a plane rotation (Q_i) to turn B_i to R_i
+
+        rhoold.copy_(rho, non_blocking=True)
+        _sym_ortho(alphahat, beta, out=(c, s, rho))
+        thetanew = torch.mul(s, alpha)
+        torch.mul(c, alpha, out=alphabar)
+
+        # Use a plane rotation (Qbar_i) to turn R_i^T to R_i^bar
+
+        rhobarold.copy_(rhobar, non_blocking=True)
+        zetaold.copy_(zeta, non_blocking=True)
+        thetabar = sbar * rho
+        rhotemp = cbar * rho
+        _sym_ortho(cbar * rho, thetanew, out=(cbar, sbar, rhobar))
+        torch.mul(cbar, zetabar, out=zeta)
+        zetabar.mul_(-sbar)
+
+        # Update h, h_hat, x.
+
+        hbar.mul_(-thetabar * rho).div_(rhoold * rhobarold)
+        hbar.add_(h)
+        x.addcdiv_(zeta * hbar, rho * rhobar)
+        h.mul_(-thetanew).div_(rho)
+        h.add_(v)
+
+        # Estimate of ||r||.
+
+        # Apply rotation Qhat_{k,2k+1}.
+        torch.mul(chat, betadd, out=betaacute)
+        torch.mul(-shat, betadd, out=betacheck)
+
+        # Apply rotation Q_{k,k+1}.
+        torch.mul(c, betaacute, out=betahat)
+        torch.mul(-s, betaacute, out=betadd)
+
+        # Apply rotation Qtilde_{k-1}.
+        # betad = betad_{k-1} here.
+
+        thetatildeold.copy_(thetatilde, non_blocking=True)
+        _sym_ortho(rhodold, thetabar, out=(ctildeold, stildeold, rhotildeold))
+        torch.mul(stildeold, rhobar, out=thetatilde)
+        torch.mul(ctildeold, rhobar, out=rhodold)
+        betad.mul_(-stildeold).addcmul_(ctildeold, betahat)
+
+        # betad   = betad_k here.
+        # rhodold = rhod_k  here.
+
+        tautildeold.mul_(-thetatildeold).add_(zetaold).div_(rhotildeold)
+        torch.div(zeta - thetatilde * tautildeold, rhodold, out=taud)
+        d.addcmul_(betacheck, betacheck)
+        torch.sqrt(d + (betad - taud).square() + betadd.square(), out=normr)
+
+        # Estimate ||A||.
+        normA2.addcmul_(beta, beta)
+        torch.sqrt(normA2, out=normA)
+        normA2.addcmul_(alpha, alpha)
+
+        # Estimate cond(A).
+        torch.max(maxrbar, rhobarold, out=maxrbar)
+        if itn > 1:
+            torch.min(minrbar, rhobarold, out=minrbar)
+
+        # ------- Test for convergence --------
+
+        # if itn % 10 == 0:
+        if True:
+
+            # Compute norms for convergence testing.
+            torch.abs(zetabar, out=normar)
+            torch.norm(x, out=normx)
+            torch.div(torch.max(maxrbar, rhotemp), torch.min(minrbar, rhotemp), out=condA)
+
+            # Now use these norms to estimate certain other quantities,
+            # some of which will be small near a solution.
+            test1 = normr / normb
+            test2 = normar / (normA * normr + eps)
+            test3 = 1 / (condA + eps)
+            t1 = test1 / (1 + normA * normx / normb)
+            rtol = btol + atol * normA * normx / normb
+
+            # The first 3 tests guard against extremely small values of
+            # atol, btol or ctol.  (The user may have set any or all of
+            # the parameters atol, btol, conlim  to 0.)
+            # The effect is equivalent to the normAl tests using
+            # atol = eps,  btol = eps,  conlim = 1/eps.
+
+            # The second 3 tests allow for tolerances set by the user.
+
+            stop = (
+                (1 + test3 <= 1)
+                | (1 + test2 <= 1)
+                | (1 + t1 <= 1)
+                | (test3 <= ctol)
+                | (test2 <= atol)
+                | (test1 <= rtol)
+            )
+
+            if stop:
+                break
+
+    return x, itn


Test patch:
diff --git a/tests/test_lsmr.py b/tests/test_lsmr.py
new file mode 100644
index 0000000..7481a95
--- /dev/null
+++ b/tests/test_lsmr.py
@@ -0,0 +1,315 @@
+"""
+Code imported from scipy
+
+Copyright (C) 2010 David Fong and Michael Saunders
+Distributed under the same license as SciPy
+
+Testing Code for LSMR.
+
+03 Jun 2010: First version release with lsmr.py
+
+David Chin-lung Fong            clfong@stanford.edu
+Institute for Computational and Mathematical Engineering
+Stanford University
+
+Michael Saunders                saunders@stanford.edu
+Systems Optimization Laboratory
+Dept of MS&E, Stanford University.
+
+"""
+
+import torch
+import unittest
+
+from torchsparsegradutils.utils import lsmr
+
+# from numpy import array, arange, eye, zeros, ones, sqrt, transpose, hstack
+# from numpy.linalg import norm
+# from numpy.testing import assert_allclose
+# import pytest
+# from scipy.sparse import coo_matrix
+# from scipy.sparse.linalg._interface import aslinearoperator
+# from scipy.sparse.linalg import lsmr
+# from .test_lsqr import G, b
+
+
+def _gettestproblem(dtype, device):
+    # Set up a test problem
+    n = 35
+    G = torch.eye(n, dtype=dtype, device=device)
+
+    for jj in range(5):
+        gg = torch.randn(n, dtype=dtype, device=device)
+        hh = torch.outer(gg, gg)
+        G += hh
+        # G += (hh + hh.reshape(-1,1)) * 0.5
+        # G += torch.randn(n, dtype=dtype, device=device) * torch.randn(n, dtype=dtype, device=device)
+
+    b = torch.randn(n, dtype=dtype, device=device)
+
+    return G, b
+
+
+class TestLSMR(unittest.TestCase):
+    def setUp(self):
+        # The device can be specialised by a daughter class
+        if not hasattr(self, "device"):
+            self.device = torch.device("cpu")
+        self.n = 10
+        self.m = 10
+        self.dtype = torch.float64
+        self.cdtype = torch.complex128
+
+    def assertCompatibleSystem(self, A, xtrue):
+        Afun = A.matmul
+        b = Afun(xtrue)
+        x = lsmr(A, b)[0]
+        # print("A",A)
+        # print("b",b)
+        # print("x",x)
+        # print("xtrue",xtrue)
+        self.assertTrue(torch.allclose(x, xtrue, atol=1e-3, rtol=1e-4))
+
+    def testIdentityACase1(self):
+        A = torch.eye(self.n, dtype=self.dtype, device=self.device)
+        xtrue = torch.zeros((self.n, 1), dtype=self.dtype, device=self.device)
+        self.assertCompatibleSystem(A, xtrue)
+
+    def testIdentityACase2(self):
+        A = torch.eye(self.n, dtype=self.dtype, device=self.device)
+        xtrue = torch.ones((self.n, 1), dtype=self.dtype, device=self.device)
+        self.assertCompatibleSystem(A, xtrue)
+
+    def testIdentityACase3(self):
+        A = torch.eye(self.n, dtype=self.dtype, device=self.device)
+        xtrue = torch.t(torch.arange(self.n, 0, -1, dtype=self.dtype, device=self.device))
+        self.assertCompatibleSystem(A, xtrue)
+
+    def testBidiagonalA(self):
+        A = lowerBidiagonalMatrix(20, self.n, self.dtype, self.device)
+        xtrue = torch.t(torch.arange(self.n, 0, -1, dtype=self.dtype, device=self.device))
+        self.assertCompatibleSystem(A, xtrue)
+
+    def testScalarB(self):
+        A = torch.tensor([[1.0, 2.0]], dtype=self.dtype, device=self.device)
+        b = torch.tensor([3.0], dtype=self.dtype, device=self.device)
+        x = lsmr(A, b)[0]
+        self.assertTrue(torch.allclose(A.matmul(x), b, atol=1e-3, rtol=1e-4))
+
+    def testComplexX(self):
+        A = torch.eye(self.n, dtype=self.cdtype, device=self.device)
+        xtrue = torch.t(
+            torch.arange(self.n, 0, -1, dtype=self.dtype, device=self.device).to(dtype=self.cdtype) * (1 + 1j)
+        )
+        self.assertCompatibleSystem(A, xtrue)
+
+    def testComplexX0(self):
+        A = 4 * torch.eye(self.n, dtype=self.cdtype, device=self.device) + torch.ones(
+            (self.n, self.n), dtype=self.cdtype, device=self.device
+        )
+        xtrue = torch.t(torch.arange(self.n, 0, -1, dtype=self.dtype, device=self.device).to(dtype=self.cdtype))
+        b = A.matmul(xtrue)
+        x0 = torch.zeros(self.n, dtype=self.cdtype, device=self.device)
+        x = lsmr(A, b, x0=x0)[0]
+        self.assertTrue(torch.allclose(x, xtrue, atol=1e-3, rtol=1e-4))
+
+    def testComplexA(self):
+        A = 4 * torch.eye(self.n, dtype=self.cdtype, device=self.device) + 1j * torch.ones(
+            (self.n, self.n), dtype=self.cdtype, device=self.device
+        )
+        xtrue = torch.t(torch.arange(self.n, 0, -1, dtype=self.dtype, device=self.device).to(dtype=self.cdtype))
+        self.assertCompatibleSystem(A, xtrue)
+
+    def testComplexB(self):
+        A = 4 * torch.eye(self.n, dtype=self.cdtype, device=self.device) + torch.ones(
+            (self.n, self.n), dtype=self.cdtype, device=self.device
+        )
+        xtrue = torch.t(
+            torch.arange(self.n, 0, -1, dtype=self.dtype, device=self.device).to(dtype=self.cdtype) * (1 + 1j)
+        )
+        b = A.matmul(xtrue)
+        x = lsmr(A, b)[0]
+        self.assertTrue(torch.allclose(x, xtrue, atol=1e-3, rtol=1e-4))
+
+    def testColumnB(self):
+        A = torch.eye(self.n, dtype=self.dtype, device=self.device)
+        b = torch.ones((self.n, 1), dtype=self.dtype, device=self.device)
+        x = lsmr(A, b)[0]
+        self.assertTrue(torch.allclose(A.matmul(x), b, atol=1e-3, rtol=1e-4))
+
+    def testInitialization(self):
+        G, b = _gettestproblem(self.dtype, self.device)
+        # Test that the default setting is not modified
+        # x_ref, _, itn_ref, normr_ref, *_ = lsmr(G, b)
+        x_ref = lsmr(G, b)[0]
+        self.assertTrue(torch.allclose(G @ x_ref, b, atol=1e-3, rtol=1e-4))
+
+        # Test passing zeros yields similiar result
+        x0 = torch.zeros(b.shape, dtype=self.dtype, device=self.device)
+        x = lsmr(G, b, x0=x0)[0]
+        self.assertTrue(torch.allclose(x, x_ref, atol=1e-3, rtol=1e-4))
+
+        # Test warm-start with single iteration
+        x0 = lsmr(G, b, maxiter=1)[0]
+
+        # x, _, itn, normr, *_ = lsmr(G, b, x0=x0)
+        x = lsmr(G, b, x0=x0)[0]
+        self.assertTrue(torch.allclose(G @ x, b, atol=1e-3, rtol=1e-4))
+
+        # NOTE(gh-12139): This doesn't always converge to the same value as
+        # ref because error estimates will be slightly different when calculated
+        # from zeros vs x0 as a result only compare norm and itn (not x).
+
+        # x generally converges 1 iteration faster because it started at x0.
+        # itn == itn_ref means that lsmr(x0) took an extra iteration see above.
+        # -1 is technically possible but is rare (1 in 100000) so it's more
+        # likely to be an error elsewhere.
+        # assert itn - itn_ref in (0, 1)
+
+        # If an extra iteration is performed normr may be 0, while normr_ref
+        # may be much larger.
+        # assert normr < normr_ref * (1 + 1e-6)
+
+    def testVerbose(self):
+        lsmrtest(20, 10, 0, self.dtype, self.device)
+
+
+class TestLSMRCUDA(TestLSMR):
+    """Override superclass setUp to run on GPU"""
+
+    def setUp(self) -> None:
+        if not torch.cuda.is_available():
+            self.skipTest(f"Skipping {self.__class__.__name__} since CUDA is not available")
+        self.device = torch.device("cuda")
+        super().setUp()
+
+
+class TestLSMRReturns(unittest.TestCase):
+    def setUp(self):
+        # The device can be specialised by a daughter class
+        if not hasattr(self, "device"):
+            self.device = torch.device("cpu")
+        self.dtype = torch.float64
+        self.n = 10
+        self.A = lowerBidiagonalMatrix(20, self.n, self.dtype, self.device)
+        self.xtrue = torch.t(torch.arange(self.n, 0, -1, dtype=self.dtype, device=self.device))
+        self.Afun = self.A.matmul
+        self.Arfun = self.A.T.matmul
+        self.b = self.Afun(self.xtrue)
+        self.x0 = torch.ones(self.n, dtype=self.dtype, device=self.device)
+        self.x00 = self.x0.clone()
+        self.returnValues = lsmr(self.A, self.b)
+        self.returnValuesX0 = lsmr(self.A, self.b, x0=self.x0)
+
+    def test_unchanged_x0(self):
+        # x, istop, itn, normr, normar, normA, condA, normx = self.returnValuesX0
+        x = self.returnValuesX0[0]
+        self.assertTrue(torch.allclose(self.x00, self.x0, atol=1e-3, rtol=1e-4))
+
+    def testNormr(self):
+        # x, istop, itn, normr, normar, normA, condA, normx = self.returnValues
+        x = self.returnValues[0]
+        self.assertTrue(torch.allclose(self.Afun(x), self.b, atol=1e-3, rtol=1e-4))
+
+    def testNormar(self):
+        # x, istop, itn, normr, normar, normA, condA, normx = self.returnValues
+        x = self.returnValuesX0[0]
+        self.assertTrue(
+            torch.allclose(
+                self.Arfun(self.b - self.Afun(x)),
+                torch.zeros(self.n, dtype=self.dtype, device=self.device),
+                atol=1e-3,
+                rtol=1e-4,
+            )
+        )
+
+    # def testNormx(self):
+    #    x, istop, itn, normr, normar, normA, condA, normx = self.returnValues
+    #    assert norm(x) == pytest.approx(normx)
+
+
+class TestLSMRReturnsCUDA(TestLSMRReturns):
+    """Override superclass setUp to run on GPU"""
+
+    def setUp(self) -> None:
+        if not torch.cuda.is_available():
+            self.skipTest(f"Skipping {self.__class__.__name__} since CUDA is not available")
+        self.device = torch.device("cuda")
+        super().setUp()
+
+
+def lowerBidiagonalMatrix(m, n, dtype, device):
+    # This is a simple example for testing LSMR.
+    # It uses the leading m*n submatrix from
+    # A = [ 1
+    #       1 2
+    #         2 3
+    #           3 4
+    #             ...
+    #               n ]
+    # suitably padded by zeros.
+    #
+    # 04 Jun 2010: First version for distribution with lsmr.py
+    if m <= n:
+        row = torch.hstack((torch.arange(m, dtype=int, device=device), torch.arange(1, m, dtype=int, device=device)))
+        col = torch.hstack((torch.arange(m, dtype=int, device=device), torch.arange(m - 1, dtype=int, device=device)))
+        idx = torch.vstack((row, col))
+        data = torch.hstack(
+            (torch.arange(1, m + 1, dtype=dtype, device=device), torch.arange(1, m, dtype=dtype, device=device))
+        )
+        return torch.sparse_coo_tensor(idx, data, size=(m, n))
+    else:
+        row = torch.hstack(
+            (torch.arange(n, dtype=int, device=device), torch.arange(1, n + 1, dtype=int, device=device))
+        )
+        col = torch.hstack((torch.arange(n, dtype=int, device=device), torch.arange(n, dtype=int, device=device)))
+        idx = torch.vstack((row, col))
+        data = torch.hstack(
+            (torch.arange(1, n + 1, dtype=dtype, device=device), torch.arange(1, n + 1, dtype=dtype, device=device))
+        )
+        return torch.sparse_coo_tensor(idx, data, size=(m, n))
+
+
+def lsmrtest(m, n, damp, dtype, device):
+    """Verbose testing of lsmr"""
+
+    A = lowerBidiagonalMatrix(m, n, dtype, device)
+    xtrue = torch.arange(n, 0, -1, dtype=dtype, device=device)
+    Afun = A.matmul
+
+    b = Afun(xtrue)
+
+    atol = 1.0e-7
+    btol = 1.0e-7
+    conlim = 1.0e10
+    itnlim = 10 * n
+    show = 1
+
+    # x, istop, itn, normr, normar, norma, conda, normx \
+    #  = lsmr(A, b, damp, atol, btol, conlim, itnlim, show)
+    x = lsmr(A, b, damp=damp, atol=atol, btol=btol, conlim=conlim, maxiter=itnlim)[0]
+
+    j1 = min(n, 5)
+    j2 = max(n - 4, 1)
+    print(" ")
+    print("First elements of x:")
+    str = ["%10.4f" % (xi) for xi in x[0:j1]]
+    print("".join(str))
+    print(" ")
+    print("Last  elements of x:")
+    str = ["%10.4f" % (xi) for xi in x[j2 - 1 :]]
+    print("".join(str))
+
+    r = b - Afun(x)
+    r2 = torch.sqrt(torch.norm(r) ** 2 + (damp * torch.norm(x)) ** 2)
+    print(" ")
+    # str = 'normr (est.)  %17.10e' % (normr)
+    str2 = "normr (true)  %17.10e" % (r2)
+    # print(str)
+    print(str2)
+    print(" ")
+
+
+if __name__ == "__main__":
+    # lsmrtest(20,10,0)
+    unittest.main()