Evaluate f_info for real in Gauss-Newton init instead of zero-filling

[Sumu004]

Thanks @NiallOswald for the root-cause tracing in #229 — this implements what you diagnosed.

AbstractGaussNewton.init filled state.f_info with zeros via tree_full_like, which also zeroes any data closed over inside f_info.jac (the jax.linearize capture). ClassicalTrustRegion.predict_reduction calls f_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 real nan that 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/InverseLevenbergMarquardt both inherit the same init and are.

Fix: call _make_f_info directly in init instead of faking its shape with zeros — exactly what step already 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

[patrick-kidger]

Thank you for the fix! (And gosh, sorry for the delay getting to this - both your fix and the original issue.)

So I think this might still actually have a slight bug: we'd never update the linearisation point from the initial value used.

This is because of these lines:

optimistix/optimistix/_solver/gauss_newton.py

Lines 254 to 259 in 6415ab1

	# We have a jaxpr in `f_info.jac`, which are compared by identity. Here we
	# arrange to use the same one so that downstream equality checks (e.g. in the
	# `filter_cond` below)
	dynamic = eqx.filter(f_eval_info.jac, eqx.is_array)
	static = eqx.filter(state.f_info.jac, eqx.is_array, inverse=True)
	jac = eqx.combine(dynamic, static)

which always keep the "old" jaxpr (as noted in the comment!) with the mistake being it does not distinguish between a Jaxpr and a ClosedJaxpr. Thus it preserves the entire ClosedJaxpr, including the constants.

I've not tracked down exactly what is holding the ClosedJaxpr - probably the correct fix is to either update it to use a Jaxpr instead, or to adjust the location to also handle ClosedJaxprs (or for robustness, both!)

WDYT?