Check finite input (#130)

* check for non-finite inputs if solution has NaN or inf

* remove private import

---------

Co-authored-by: Johanna Haffner <johanna.haffner@bsse.ethz.ch>

Author: 2 people

lineax/_solution.py

@@ -40,6 +40,15 @@
 """.strip()
 
 
+_nonfinite_msg = """
+The linear solver received non-finite (NaN or inf) input and cannot determine a 
+solution. 
+
+This means that you have a bug upstream of lineax and should check the inputs to 
+`lineax.linear_solve` for non-finite values.
+""".strip()
+
+
 class RESULTS(eqxi.Enumeration):
     successful = ""
     max_steps_reached = (
@@ -55,6 +64,7 @@ class RESULTS(eqxi.Enumeration):
         "A stagnation in an iterative linear solve has occurred. Try increasing "
         "`stagnation_iters` or `restart`."
     )
+    nonfinite_input = _nonfinite_msg
 
 
 class Solution(eqx.Module, strict=True):

lineax/_solve.py

@@ -90,16 +90,26 @@ def _linear_solve_impl(_, state, vector, options, solver, throw, *, check_closur
             out, name="lineax.linear_solve with respect to a closed-over value"
         )
     solution, result, stats = out
-    has_nonfinites = jnp.any(
+    has_nonfinite_output = jnp.any(
         jnp.stack(
             [jnp.any(jnp.invert(jnp.isfinite(x))) for x in jtu.tree_leaves(solution)]
         )
     )
     result = RESULTS.where(
-        (result == RESULTS.successful) & has_nonfinites,
+        (result == RESULTS.successful) & has_nonfinite_output,
         RESULTS.singular,
         result,
     )
+    has_nonfinite_input = jnp.any(
+        jnp.stack(
+            [jnp.any(jnp.invert(jnp.isfinite(x))) for x in jtu.tree_leaves(vector)]
+        )
+    )
+    result = RESULTS.where(
+        (result == RESULTS.singular) & has_nonfinite_input,
+        RESULTS.nonfinite_input,
+        result,
+    )
     if throw:
         solution, result, stats = result.error_if(
             (solution, result, stats),

tests/test_solve.py

@@ -191,3 +191,18 @@ def test_iterative_solver_max_steps_only(solver):
     rhs = jax.random.normal(jax.random.key(0), (SIZE,))
 
     lx.linear_solve(poisson_operator, rhs, solver)
+
+
+def test_nonfinite_input():
+    operator = lx.DiagonalLinearOperator((1.0, 1.0))
+    vector = (1.0, jnp.inf)
+    sol = lx.linear_solve(operator, vector, throw=False)
+    assert sol.result == lx.RESULTS.nonfinite_input
+
+    vector = (1.0, jnp.nan)
+    sol = lx.linear_solve(operator, vector, throw=False)
+    assert sol.result == lx.RESULTS.nonfinite_input
+
+    vector = (jnp.nan, jnp.inf)
+    sol = lx.linear_solve(operator, vector, throw=False)
+    assert sol.result == lx.RESULTS.nonfinite_input