Merge pull request #176 from johannahaffner/lsmr-init-overflow

Decrease default value to prevent overflow in 32-bit.

Author: johannahaffner

lineax/_solver/lsmr.py

@@ -42,6 +42,7 @@
 from equinox.internal import ω
 from jaxtyping import Array, PyTree
 
+from .._misc import complex_to_real_dtype
 from .._norm import two_norm
 from .._operator import AbstractLinearOperator, conj
 from .._solution import RESULTS
@@ -108,23 +109,30 @@ def compute(
             and self.rtol == 0
         )
 
+        dtype = jnp.result_type(
+            *jtu.tree_leaves(vector),
+            *jtu.tree_leaves(x),
+            *jtu.tree_leaves(operator.in_structure()),
+        )
+
         m, n = operator.out_size(), operator.in_size()
         # number of singular values
         min_dim = min([m, n])
         if self.max_steps is None:
-            max_steps = min_dim * 10  # for consistency with other iterative solvers
+            # Set max_steps based on the minimum dimension + avoid numerical overflows
+            # https://github.com/patrick-kidger/lineax/issues/175
+            # https://github.com/patrick-kidger/lineax/issues/177
+            int_dtype = jnp.dtype(f"int{complex_to_real_dtype(dtype).itemsize * 8}")
+            if min_dim > (jnp.iinfo(int_dtype).max / 10):
+                max_steps = jnp.iinfo(int_dtype).max
+            else:
+                max_steps = min_dim * 10  # for consistency with other iterative solvers
         else:
             max_steps = self.max_steps
 
         if x is None:
             x = jtu.tree_map(jnp.zeros_like, operator.in_structure())
 
-        dtype = jnp.result_type(
-            *jtu.tree_leaves(vector),
-            *jtu.tree_leaves(x),
-            *jtu.tree_leaves(operator.in_structure()),
-        )
-
         b = vector
         u = (ω(b) - ω(operator.mv(x))).ω
         normb = self.norm(b)
@@ -178,7 +186,7 @@ def beta_zero(beta, u):
             # variables for estimation of ||A|| and cond(A)
             normA2=alpha**2,
             maxrbar=0.0,
-            minrbar=1e100,
+            minrbar=jnp.finfo(dtype).max,
             condA=1.0,
             # variables for use in stopping rules
             istop=0,