fix: prevent BiCGStab producing NaN on exact/near-exact intermediate convergence

[jpbrodrick89]

The bug

BiCGStab.compute's body_fun computes omega_new = <s, t> / <t, t>, where s is the half-step residual and t = A(M(s)). This leads to a divide by zero error whenever s (and hence t) is already (numerically) the zero vector and happens whenever r0 lands in an invariant subspace of the preconditioned operator (e.g. b is exactly an eigenvector of A). Consequently, BiCGStab in effect solves the system in fewer steps than the recurrence "expects".

This is caught by RESULTS.nonfinite_input meaning we fail on problems that are perfectly well-posed.

Industry standard

scipy.sparse.linalg.bicgstab guards exactly this case (if norm(s) < atol: x += alpha*phat; return x, 0), and jax.scipy.sparse.linalg.bicgstab copies the same guard (exit_early). A second, related cascade one iteration later (alpha_new = rho_new / <r0, v_new>, both tiny once r is already converged) is also guarded here for the same reason.

MWE

A = jnp.diag(jnp.array([2.0, -3.0, 5.0, 7.0]))
b = jnp.array([1.0, 0.0, 0.0, 0.0])  # exact eigenvector of A, eigenvalue 2.0
lx.linear_solve(lx.MatrixLinearOperator(A), b, solver=lx.BiCGStab(rtol=1e-8, atol=1e-8))
# before: value=[nan, nan, nan, nan], result=nonfinite output
# after:  value=[0.5, 0., 0., 0.], result=successful

Fix

Mirror scipy and jax using jnp.where:

• Guard omega_new's division when the half-step residual s is already within the solver's own convergence tolerance (r_already_converged), reporting omega_new = 1 (not 0, so as not to spuriously trip breakdown_occurred, which treats omega == 0 as true stagnation).
• Guard alpha_new's division the same way when the incoming residual r is already converged and force a genuine no-op step (diff = 0, r unchanged) so not_converged correctly detects quiescence on the next loop check.