Test failure with 0.7.1: condition number too high for LSMR solver; test passes on 0.7.0

[johannahaffner]

I ran the Lineax test suite for 0.7.1, with Equinox from its current main branch.

We now get an EquinoxRuntimeError for LSMR.

tests/test_vmap_jvp.py:129: EquinoxRuntimeError
=========================== short test summary info ============================
FAILED tests/test_vmap_jvp.py::test_vmap_jvp[float64-construct_singular_matrix-True-make_matrix_operator-solver20-tags20-True] - EquinoxRuntimeError: Above is the stack outside of JIT. Below is the stack ...

The Lineax error message is

equinox.EquinoxRuntimeError: Condition number of A seems to be larger than `conlim`.

This test passes on 0.7.0.
I have not yet investigated this error, just opening this issue so it can be tracked.

[f0uriest]

I see the same issue with jax==0.7.1 but seems fine on 0.7.2 or 0.8.0, and I don't see anything in the changelog that would explain the difference. I don't think jax guarantees bit for bit identical rng output between versions, so it's possible we just got unlucky with the randomly generated matrix?

[johannahaffner]

so it's possible we just got unlucky with the randomly generated matrix?

I think so too.