fix: prevent LM damping underflow

docs/callbacks.md

@@ -44,7 +44,10 @@ and `xtol` cannot fire at step zero.
 
 Repeated rejections multiply the damping by `damping_increase` without bound,
 which can overflow in float32. The constructor's `max_damping` clamps the
-damping from above; leave it `None` for uncapped classic behavior.
+damping from above; leave it `None` for uncapped classic behavior. Accepted
+steps cannot underflow damping to zero: `min_damping=None` uses
+`jnp.finfo(residual.dtype).tiny`, while an explicit value selects a larger
+absolute floor.
 
 Status codes are integer constants:
 
@@ -82,10 +85,10 @@ to second-guess a callback's explicit replacement.
 ### Resettable Hyperparameters
 
 `solve()` populates `lm_state.hyper` with an `LMHyperparams` of traced
-per-step values: `damping_decrease`, `damping_increase`, `max_damping`,
-`geodesic_acceptance_ratio`, `iterative_tol`, `iterative_atol`, and
-`iterative_maxiter`. Because they ride in the lm_state, a
-callback can reset any of them mid-solve — exactly like a damping reset:
+per-step values: `damping_decrease`, `damping_increase`, `min_damping`,
+`max_damping`, `geodesic_acceptance_ratio`, `iterative_tol`, `iterative_atol`,
+and `iterative_maxiter`. Because they ride in the lm_state, a callback can reset
+any of them mid-solve — exactly like a damping reset:
 
 ```python
 new_hyper = dataclasses.replace(

docs/index.md

@@ -140,7 +140,9 @@ so the natural parameter metric is \(M=K\).
 
 The actual nonlinear step is accepted only if the unregularized residual sum of
 squares decreases. On acceptance, damping is multiplied by `damping_decrease`;
-on rejection, it is multiplied by `damping_increase`.
+on rejection, it is multiplied by `damping_increase`. `min_damping=None`
+clamps it to the smallest positive normal value of the residual dtype, avoiding
+backend flush-to-zero behavior without imposing a larger regularization floor.
 
 Near the interpolation threshold, small-damping LM becomes metric
 Gauss-Newton, whose step is the minimum-\(M\)-norm solution of the linearized

docs/tuning_guide.md

@@ -192,6 +192,11 @@ few steps. Try them when you see specific signatures:
   descent).
 - Long rejection storms in float32 → set `max_damping` (~`1e6`) so damping
   cannot overflow.
+- Literal damping underflow → the default `min_damping=None` already resolves
+  to the residual dtype's smallest positive normal value. This representation
+  floor prevents flush-to-zero but does not keep the damping numerically active
+  in an ill-conditioned system. A standard scale-aware floor is on the order of
+  `eps * operator_scale`; pass that larger absolute `min_damping` explicitly.
 - Accept/reject oscillation → bring `damping_decrease`/`damping_increase`
   closer to 1 (e.g. 0.7 / 2.0) for smoother adaptation.
 - All steps accepted but progress is slow → lower `init_damping` or decrease

src/nlls_gram/gram_lm.py

@@ -64,19 +64,31 @@ class LMHyperparams:
 
     damping_decrease: jax.Array
     damping_increase: jax.Array
+    min_damping: jax.Array
     max_damping: jax.Array | None
     geodesic_acceptance_ratio: jax.Array
     iterative_tol: jax.Array
     iterative_atol: jax.Array
     iterative_maxiter: jax.Array | None
 
 
+def _damping_floor(min_damping, dtype):
+    if dtype is None:
+        seed = 0.0 if min_damping is None else min_damping
+        dtype = jnp.asarray(seed).dtype
+    dtype_floor = jnp.asarray(jnp.finfo(dtype).tiny, dtype=dtype)
+    if min_damping is None:
+        return dtype_floor
+    return jnp.maximum(jnp.asarray(min_damping, dtype=dtype), dtype_floor)
+
+
 def _cast_hyper(hyper, dtype):
     if hyper is None:
         return None
     return LMHyperparams(
         jnp.asarray(hyper.damping_decrease, dtype=dtype),
         jnp.asarray(hyper.damping_increase, dtype=dtype),
+        _damping_floor(hyper.min_damping, dtype),
         None
         if hyper.max_damping is None
         else jnp.asarray(hyper.max_damping, dtype=dtype),
@@ -984,6 +996,15 @@ class LevenbergMarquardt:
     Take higher-order derivatives of ``solve`` with a fixed metric. See
     :class:`MetricFactory`.
 
+    ``min_damping`` is the absolute lower bound applied before every linear
+    solve and after every damping update. ``None`` (the default) resolves to
+    ``jnp.finfo(residual.dtype).tiny``, the smallest positive normal value, so
+    damping cannot enter a backend's flush-to-zero range. This is only an
+    underflow floor. A standard conditioning-oriented minimum is instead on
+    the order of machine epsilon times a representative scale of ``B'B`` (or
+    ``B B'``); pass that larger absolute value explicitly when the damping must
+    remain numerically effective in the linear system.
+
     ``solve(...).x`` has a custom implicit AD rule with respect to ``p``,
     relinearized at the returned solution. ``ad_solver`` explicitly selects
     one method from ``{"direct", "svd", "qr", "augmented_qr", "gram_cg",
@@ -1104,6 +1125,7 @@ def __init__(
         init_damping=1e-3,
         damping_decrease=0.5,
         damping_increase=4.0,
+        min_damping=None,
         max_damping=None,
         linear_solver="auto",
         jacobian_mode="auto",
@@ -1150,6 +1172,10 @@ def __init__(
             raise ValueError("damping_decrease must be positive")
         if damping_increase <= 0:
             raise ValueError("damping_increase must be positive")
+        if min_damping is not None and min_damping <= 0:
+            raise ValueError("min_damping must be positive or None")
+        if min_damping is not None and min_damping > init_damping:
+            raise ValueError("min_damping must not exceed init_damping")
         if max_damping is not None and max_damping < init_damping:
             raise ValueError("max_damping must be at least init_damping")
         if iterative_tol < 0:
@@ -1365,6 +1391,7 @@ def __init__(
         self.init_damping = init_damping
         self.damping_decrease = damping_decrease
         self.damping_increase = damping_increase
+        self.min_damping = min_damping
         self.max_damping = max_damping
         self.linear_solver = linear_solver
         self.jacobian_mode = jacobian_mode
@@ -1445,6 +1472,7 @@ def __init__(
                 init_damping,
                 damping_decrease,
                 damping_increase,
+                min_damping,
                 max_damping,
                 linear_solver,
                 jacobian_mode,
@@ -1489,6 +1517,7 @@ def hyperparams(self, dtype=None):
         return LMHyperparams(
             jnp.asarray(self.damping_decrease, dtype=dtype),
             jnp.asarray(self.damping_increase, dtype=dtype),
+            _damping_floor(self.min_damping, dtype),
             None
             if self.max_damping is None
             else jnp.asarray(self.max_damping, dtype=dtype),
@@ -1508,7 +1537,10 @@ def init(self, x0, args=None, *, p=None):
         # populates it for its callbacks.
         self._check_residual_args(args, p)
         residual, aux = self._residual_and_aux(x0, args, p)
-        damping = jnp.asarray(self.init_damping, dtype=residual.dtype)
+        min_damping = _damping_floor(self.min_damping, residual.dtype)
+        damping = jnp.maximum(
+            jnp.asarray(self.init_damping, dtype=residual.dtype), min_damping
+        )
         recycle = self._init_recycle_state(residual)
         precond, precond_valid = self._init_precond(x0, args, p, aux)
         metric_state, metric_valid = self._init_metric_state(x0, args, p, aux)
@@ -1728,7 +1760,6 @@ def reuse_resid_and_jt(_):
             )
         else:
             resid, Jt, aux = self._dense_resid_jt_aux(residual_flat, theta)
-        damping = jnp.asarray(lm_state.damping, dtype=resid.dtype)
         # Traced hyperparameters from the lm_state when present (resettable by
         # solve callbacks); the None fallback compiles to the same constants
         # as reading the constructor values directly.
@@ -1739,6 +1770,10 @@ def reuse_resid_and_jt(_):
         )
         damping_decrease = jnp.asarray(hyper.damping_decrease, dtype=resid.dtype)
         damping_increase = jnp.asarray(hyper.damping_increase, dtype=resid.dtype)
+        min_damping = _damping_floor(hyper.min_damping, resid.dtype)
+        damping = jnp.maximum(
+            jnp.asarray(lm_state.damping, dtype=resid.dtype), min_damping
+        )
 
         # Trace-time shape resolution of the default solver: shapes are concrete
         # while tracing, so this is a plain Python branch (never a lax.cond).
@@ -2231,8 +2266,12 @@ def accelerated_loss(_):
         new_damping = damping * damping_factor
         if hyper.max_damping is not None:
             new_damping = jnp.minimum(
-                new_damping, jnp.asarray(hyper.max_damping, dtype=resid.dtype)
+                new_damping,
+                jnp.maximum(
+                    jnp.asarray(hyper.max_damping, dtype=resid.dtype), min_damping
+                ),
             )
+        new_damping = jnp.maximum(new_damping, min_damping)
         loss = jnp.where(improved, loss_candidate, loss_old)
         # New recycle state: the velocity solve's harvested basis and the (stop-
         # gradient'd) dual solutions become warm starts for the next step. Threads

tests/test_float64_subprocess.py

@@ -1552,3 +1552,39 @@ def residual(theta):
         capture_output=True,
         text=True,
     )
+
+
+def test_float64_default_min_damping_uses_float64_normal_floor():
+    script = r"""
+import jax
+jax.config.update("jax_enable_x64", True)
+import jax.numpy as jnp
+
+from nlls_gram import LevenbergMarquardt, LMState
+
+
+def residual(theta):
+    return theta
+
+
+solver = LevenbergMarquardt(
+    residual,
+    cache_jacobian=False,
+    geodesic_acceleration=False,
+)
+state = LMState(jnp.asarray(0.0, dtype=jnp.float64))
+x, state, info = solver.update(jnp.ones(1, dtype=jnp.float64), state)
+floor = jnp.asarray(jnp.finfo(jnp.float64).tiny, dtype=jnp.float64)
+
+assert info.accepted
+assert jnp.all(jnp.isfinite(x))
+assert state.damping.dtype == jnp.float64
+assert state.damping == floor, (state.damping, floor)
+assert info.damping == floor, (info.damping, floor)
+"""
+    subprocess.run(
+        [sys.executable, "-c", textwrap.dedent(script)],
+        check=True,
+        capture_output=True,
+        text=True,
+    )

tests/test_gram_lm.py

@@ -234,6 +234,13 @@ def test_damping_update_factors_must_be_positive():
         LevenbergMarquardt(residual_fn, damping_increase=0.0)
 
 
+def test_min_damping_must_be_positive_and_not_exceed_initial_damping():
+    with pytest.raises(ValueError, match="min_damping must be positive or None"):
+        LevenbergMarquardt(residual_fn, min_damping=0.0)
+    with pytest.raises(ValueError, match="min_damping must not exceed"):
+        LevenbergMarquardt(residual_fn, init_damping=1e-3, min_damping=1e-2)
+
+
 def test_iterative_options_must_be_valid():
     with pytest.raises(ValueError, match="iterative_tol must be nonnegative"):
         LevenbergMarquardt(residual_fn, linear_solver="gram_cg", iterative_tol=-1.0)
@@ -361,6 +368,56 @@ def test_max_damping_below_init_damping_raises():
         LevenbergMarquardt(residual_fn, init_damping=1e-2, max_damping=1e-3)
 
 
+@pytest.mark.parametrize("jit", [False, True])
+@pytest.mark.parametrize("linear_solver", ["auto", "lsmr"])
+def test_default_min_damping_repairs_zero_and_prevents_float32_underflow(
+    jit, linear_solver
+):
+    def residual(theta):
+        return theta
+
+    solver = LevenbergMarquardt(
+        residual,
+        init_damping=1e-3,
+        linear_solver=linear_solver,
+        cache_jacobian=False,
+        geodesic_acceleration=False,
+    )
+    state = LMState(jnp.asarray(0.0, dtype=jnp.float32))
+    update = solver.update if not jit else jax.jit(lambda x, s: solver.update(x, s))
+    x, state, info = update(jnp.ones(1, dtype=jnp.float32), state)
+    floor = jnp.asarray(jnp.finfo(jnp.float32).tiny, dtype=jnp.float32)
+
+    assert bool(info.accepted)
+    assert jnp.all(jnp.isfinite(x))
+    assert state.damping == floor
+    assert info.damping == floor
+
+
+@pytest.mark.parametrize("jit", [False, True])
+def test_explicit_min_damping_is_used_by_step_and_update(jit):
+    def residual(theta):
+        return theta
+
+    min_damping = 1e-4
+    solver = LevenbergMarquardt(
+        residual,
+        init_damping=1e-3,
+        min_damping=min_damping,
+        cache_jacobian=False,
+        geodesic_acceleration=False,
+    )
+    state = LMState(jnp.asarray(0.0, dtype=jnp.float32))
+    update = solver.update if not jit else jax.jit(lambda x, s: solver.update(x, s))
+    x, state, info = update(jnp.ones(1, dtype=jnp.float32), state)
+    expected_x = min_damping / (1.0 + min_damping)
+
+    assert bool(info.accepted)
+    assert jnp.allclose(x, expected_x, rtol=1e-4, atol=1e-7)
+    assert float(state.damping) == pytest.approx(min_damping)
+    assert float(info.damping) == pytest.approx(min_damping)
+
+
 def test_metric_requirements_per_linear_solver():
     with pytest.raises(ValueError, match="metric.solve"):
         LevenbergMarquardt(
@@ -2053,6 +2110,7 @@ def build(**overrides):
     # Any static-setting change (or a different residual function) is a
     # different solver, so it cannot silently reuse the wrong compiled loop.
     assert a != build(init_damping=2e-2)
+    assert a != build(min_damping=1e-8)
     assert a != build(geodesic_acceleration=False)
     assert a != LevenbergMarquardt(
         lambda theta, args, p: theta - args, init_damping=1e-2, cache_jacobian=False
@@ -2216,6 +2274,8 @@ def test_hyperparams_typing_and_solve_population():
     assert solver.init({"a": 1.0, "b": 0.0}, (ts, ys)).hyper is None
     hyper = solver.hyperparams(jnp.float32)
     assert hyper.damping_decrease.dtype == jnp.float32
+    assert hyper.min_damping.dtype == jnp.float32
+    assert hyper.min_damping == jnp.finfo(jnp.float32).tiny
     assert hyper.iterative_maxiter.dtype == jnp.int32
     assert hyper.max_damping is None
     result = solver.solve({"a": 1.0, "b": 0.0}, (ts, ys), max_steps=2)