LSMR fails with NaN output for trivial problems

[tvercaut]

Thanks for the nice library. I wanted to try LSMR but my first attempt to use it with a trivial problem failed with NaN output.

Steps to reproduce:

import torch
import torchmin

A = torch.eye(10)
xtrue = torch.zeros((10, 1))
b = A @ xtrue

x = torchmin.lstsq.lsmr.lsmr(A,b)[0]
print(x)

which resulted in

tensor([nan, nan, nan, nan, nan, nan, nan, nan, nan, nan])

instead of 0s

[tvercaut]

Note that the same holds with a slightly less trivial but still trivial case:

import torch
import torchmin

A = torch.eye(10)
xtrue = torch.ones((10, 1))
b = A @ xtrue

x = torchmin.lstsq.lsmr.lsmr(A,b)[0]
print(x)

Comparing the source code with that of scipy points to the following issues.

normr and normar are initialiased differently and do not include the "convergence" tests at the very beginning

pytorch-minimize/torchmin/lstsq/lsmr.py

Lines 148 to 149 in 1017e97

	normar = b.new_tensor(0)
	normr = b.new_tensor(0)

instead of
https://github.com/scipy/scipy/blob/2347d9309fbbabb1d3f89d35b7a42d0d53f002b2/scipy/sparse/linalg/_isolve/lsmr.py#L290-L307

Within the main iteration loop, the convergence test is only done every 10 iterations

pytorch-minimize/torchmin/lstsq/lsmr.py

Lines 257 to 259 in 1017e97

	# ------- Test for convergence --------
	if itn % 10 == 0:

instead of every time
https://github.com/scipy/scipy/blob/2347d9309fbbabb1d3f89d35b7a42d0d53f002b2/scipy/sparse/linalg/_isolve/lsmr.py#L409

[tvercaut]

For the record, I made a few changes here:
https://github.com/cai4cai/torchsparsegradutils/blob/main/torchsparsegradutils/utils/lsmr.py