Bug in QR solver for Tikhnov regularised matrices

[packquickly]

When solving against matrices of the form $A + \lambda I$ with $\lambda$ set to the max dtype value, QR solves return nan instead of zero. This does not occur with matrices of the form $\lambda I$ and for some matrices $A$ it seems not to be an issue, but for dense $A$ it is a problem.

This case arises in the Levenberg-Marquardt algorithm in Optimistix when the step-size goes to zero.

This failure is not silent, but it should be supported.

MWE:

import jax
import jax.numpy as jnp
import jax.random as jr
import lineax as lx

jax.config.update("jax_enable_x64", True)

vector = jnp.ones(100)
# does not need to be a `MatrixLinearOperator`.
dense_operator = lx.MatrixLinearOperator(0.1 * jr.normal(jr.PRNGKey(0), (100, 100)))
dampening = jnp.finfo(jnp.asarray(0.0)).max
tikhnov_operator = dense_operator + dampening * lx.IdentityLinearOperator(
    jax.eval_shape(lambda: vector)
)
linear_sol = lx.linear_solve(tikhnov_operator, vector, solver=lx.QR(), throw=False)

assert not jnp.any(jnp.isnan(linear_sol.value))  # assertion fails

[patrick-kidger]

Wouldn't the resulting operating contain infinities? (When you add something positive to the max value.) If so I think this is expected.

[packquickly]

I don't think this is what's happening. If you replace the matrix in dense_operator with the matrix containing only its diagonal, or with an IdentityLinearOperator, then no issue arises and the result is zero. If getting inf on the diagonals was the issue then you'd expect this to still cause problems.

Also this works with SVD and is a case we do use explicitly in Optimistix, so if it is not a bug here it should be addressed there.