LSMR init overflow

[GJBoth]

Thanks for implementing the LSMR solver - I've been using it for matrix-free Levenberg-Marquardt and it seems to work really well! There's one issue with the initialisation however - one of the initial values gets initialised as 1e100 leading to overflows.. The origin of this value seems to be the scipy implementation, which runs in 64bit by default - but Jax runs in 32. It doesn't seem to affect the result of the solver, though.

Simply changing this line to:

minrbar=jnp.inf

solves the issue for me.

[johannahaffner]

This looks reasonable to me, @f0uriest and @PTNobel do you concur? Not sure if this would then require a jnp.where or some other type of safety mechanism in the body of the solver + if a smaller numerical value than 1e100 would be preferable here.

[f0uriest]

Yeah I'm 99.9% sure using jnp.inf would be fine, I don't think it needs any where etc. The only possible issue I could see is if using jax.debug_infs you might get a false positive because of the hard-coded value? Could also use jnp.finfo(x.dtype).max?

[johannahaffner]

This just reminded me of this issue: patrick-kidger/optimistix#186 - let's try to avoid using a jnp.something default value. How about we set minrbar to the highest value that does not overflow in 32-bit, so 1e38 if this is a float, and 1e9 if it is an integer.

[PTNobel]

I think any big number should be fine ...