Evaluate f_info for real in Gauss-Newton init instead of zero-filling#234
Open
Sumu004 wants to merge 1 commit into
Open
Evaluate f_info for real in Gauss-Newton init instead of zero-filling#234Sumu004 wants to merge 1 commit into
Sumu004 wants to merge 1 commit into
Conversation
…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 patrick-kidger#229
Sumu004
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Thanks @NiallOswald for the root-cause tracing in #229 — this implements what you diagnosed.
AbstractGaussNewton.initfilledstate.f_infowith zeros viatree_full_like, which also zeroes any data closed over insidef_info.jac(thejax.linearizecapture).ClassicalTrustRegion.predict_reductioncallsf_info.jac.mv(...)unconditionally on every step, including the first, so a residual function that's nan-sensitive at a zeroed closure value produces a realnanthat contaminates the JVP — even though the accept/reject decision itself is separately forced correct on the first step.GaussNewton(no trust region) is unaffected;LevenbergMarquardt/InverseLevenbergMarquardtboth inherit the sameinitand are.Fix: call
_make_f_infodirectly ininitinstead of faking its shape with zeros — exactly whatstepalready does every other iteration. Costs one extra linearization at init.I don't have a working JAX install, so I've verified this by tracing the exact call paths rather than running it — would appreciate a maintainer confirming against the real reproducer.
Fixes #229