From 527b9c95a8091cb6c4842492037e459393d54bec Mon Sep 17 00:00:00 2001 From: Sumukh Chaluvaraju Date: Tue, 7 Jul 2026 22:55:26 +0100 Subject: [PATCH] fix(gauss_newton): evaluate f_info for real in init, don't zero-fill closures AbstractGaussNewton.init filled state.f_info with a same-shaped placeholder of zeros (tree_full_like(f_info_struct, 0, ...)) rather than a real evaluation of _make_f_info. This silently zeroed out any values closed over inside f_info.jac too -- in particular, the point at which jax.linearize captured its linearization inside _make_f_info. ClassicalTrustRegion.predict_reduction calls f_info.jac.mv(y_diff) unconditionally on every call to step -- including the very first, where f_info is still this placeholder -- even though the decision it feeds into (accept) is separately forced to True via accept | first_step. If the residual function involves an operation that is nan-sensitive at zero (e.g. a division), the corrupted placeholder produces a genuine nan in predicted_reduction's value, which then contaminates the JVP even though the decision itself is unaffected by it. GaussNewton (no trust region) doesn't hit this, which is why the original report singles out LevenbergMarquardt/InverseLevenbergMarquardt specifically -- both inherit this same init with no override. Fix: call _make_f_info(fn, y, args, tags, jac) directly in init (y is already available as a parameter), matching what step already does for subsequent iterations, instead of eval_shape + zero-fill. This costs one extra linearization at init in exchange for never fabricating ill-defined closure data. Fixes #229 --- optimistix/_solver/gauss_newton.py | 16 ++++++++++-- tests/test_root_find.py | 39 ++++++++++++++++++++++++++++++ 2 files changed, 53 insertions(+), 2 deletions(-) diff --git a/optimistix/_solver/gauss_newton.py b/optimistix/_solver/gauss_newton.py index 8dcd7f3a..4d1f581b 100644 --- a/optimistix/_solver/gauss_newton.py +++ b/optimistix/_solver/gauss_newton.py @@ -224,8 +224,20 @@ def init( tags: frozenset[object], ) -> _GaussNewtonState: jac = options.get("jac", "fwd") - f_info_struct, _ = eqx.filter_eval_shape(_make_f_info, fn, y, args, tags, jac) - f_info = tree_full_like(f_info_struct, 0, allow_static=True) + # Evaluate `f_info` for real at the initial `y`, rather than filling a + # same-shaped placeholder with zeros. `step`'s very first call reads this + # placeholder unconditionally: `predict_reduction` (e.g. in + # `ClassicalTrustRegion`) calls `f_info.jac.mv(...)` regardless of + # `first_step`, even though the decision it feeds into is later forced by + # `accept | first_step`. Zero-filling silently zeroes any values closed + # over inside `f_info.jac` (e.g. `jax.linearize`'s captured linearization + # point), so any nan-sensitive operation in the user's own residual (e.g. + # a division) produces a genuine nan here -- which then contaminates the + # JVP even though the decision itself is unaffected. This costs one extra + # linearization at init, in exchange for never fabricating ill-defined + # closure data. + f_info, _ = _make_f_info(fn, y, args, tags, jac) + f_info_struct = jax.eval_shape(lambda: f_info) return _GaussNewtonState( first_step=jnp.array(True), y_eval=y, diff --git a/tests/test_root_find.py b/tests/test_root_find.py index 11a4ddc6..2fbc6088 100644 --- a/tests/test_root_find.py +++ b/tests/test_root_find.py @@ -2,6 +2,7 @@ import random import equinox as eqx +import jax import jax.numpy as jnp import jax.random as jr import jax.tree_util as jtu @@ -243,3 +244,41 @@ def f(y, _): solver = solver_cls(rtol=1e-5, atol=1e-5, cauchy_termination=False) sol = optx.root_find(f, solver, 0.0, throw=False) assert sol.result == optx.RESULTS.successful + + +def test_levenberg_marquardt_no_nan_jvp_on_first_step(): + # Regression test for https://github.com/patrick-kidger/optimistix/issues/229 + # + # `AbstractGaussNewton.init` used to fill `state.f_info` with a same-shaped + # placeholder of zeros (via `tree_full_like(f_info_struct, 0, ...)`), rather + # than a real evaluation of `_make_f_info`. This silently zeroed out any + # values closed over inside `f_info.jac` too -- in particular, the point at + # which `jax.linearize` captured its linearization inside `_make_f_info`. + # + # `ClassicalTrustRegion.predict_reduction` calls `f_info.jac.mv(y_diff)` + # unconditionally on every call to `step` -- including the very first, where + # `f_info` is still this placeholder -- even though the *decision* it feeds + # into (`accept`) is separately forced to `True` via `accept | first_step`. + # If the residual function involves an operation that is nan-sensitive at + # zero (e.g. a division), the corrupted placeholder produces a genuine nan + # in `predicted_reduction`'s value, which then contaminates the JVP even + # though the *decision* itself is unaffected by it. + # + # `GaussNewton` (no trust region) does not hit this, which is why the + # original report singles out `LevenbergMarquardt`/ + # `InverseLevenbergMarquardt` specifically. + divisor = jnp.array([5.0, 2.5]) + + def residual(y, _): + return jnp.array([0.5, 0.5]) - y / divisor + + solver = optx.LevenbergMarquardt(rtol=1e-5, atol=1e-5) + y0 = jnp.array([0.9, 0.1]) + + def solve(y0): + return optx.root_find(residual, solver, y0=y0, args=()).value + + with jax.numpy_dtype_promotion("standard"): + out, tangent = jax.jvp(solve, (y0,), (jnp.ones_like(y0),)) + assert jnp.all(jnp.isfinite(out)) + assert jnp.all(jnp.isfinite(tangent))