test: JAX finite-difference gradient correctness suite (jax_grad)#157
Merged
Conversation
Upgrades the jax_grad scripts from finiteness-only checks to autodiff-vs- central-finite-difference correctness, per the JAX gradient audit (PyAutoLabs/autolens_workspace_developer#87): - util.py: shared FD comparison helpers (per-parameter scaled steps, eager autodiff + jitted FD sweep guarded by an eager-vs-jit consistency check, explicit skip-list semantics for knowingly-approximate parameters). - imaging_lp.py: standard lp.Sersic + lp_linear.Sersic variants, FD-validated; evaluation point anchored near the simulator truth so the positive-only NNLS keeps the source component live (prior medians zero it and legitimately kill all source/mass gradients). - imaging_mge.py: MGE source FD-validated through the linear inversion. - imaging_pixelization.py (new): RectangularUniform strictly FD-validated on all parameters; RectangularAdaptDensity asserts lens-light FD-correctness plus the staircase invariance of the likelihood in mass/shear directions (autodiff zero is the correct a.e. derivative; the rank-space CDF transform is invariant under order-preserving deformations at os_pix=1). - point_source.py (new): FitPositionsSource source-plane chi-squared FD-validated incl. the magnification-via-Hessian term (eager path; forward jit remains blocked upstream). - weak.py (new): FitWeak likelihood FD-validated, plain + redshift-scaled. Also fixes the stale dataset pointers (imaging/simple, imaging/source_complex -> imaging/jax_test) that broke both pre-existing scripts on a clean checkout. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01XGDp54jKUd3kUziTF77pNu
…ling Extends imaging_pixelization.py with two production-shaped variants at over_sample_size_pixelization=4: - RectangularAdaptImage + reg.Adapt + al.AdaptImages + border relocator (the production configuration): all 14 gradients live, AD vs FD(h=1e-7) <= ~1% on mass/shear and 6 digits on lens light. - RectangularAdaptDensity: all 14 gradients live, <= ~3.3% (worst: einstein_radius). FD uses rel_step=1e-7 (below the rank-reordering scale) with a 5% tolerance that reflects micro-staircase contamination of the finite differences — FD drifts toward autodiff as h shrinks, so autodiff is the h-consistent reference. A liveness assertion guards every parameter at os_pix=4, since this is the configuration gradient-based inference will use. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01XGDp54jKUd3kUziTF77pNu
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.
Part of the JAX autodiff gradients audit — PyAutoLabs/autolens_workspace_developer#87.
Upgrades
scripts/jax_grad/from finiteness-only checks to autodiff-vs-finite-difference correctness, and extends coverage to every differentiable PyAutoLens likelihood.Scripts Changed
scripts/jax_grad/util.py(new) — shared FD comparison helpers: central differences with per-parameter scaled steps, autodiff on the eager likelihood + FD sweep on a jitted one guarded by an eager-vs-jit base-point check (anti-pure_callback-const-folding), explicit skip-list semantics.scripts/jax_grad/imaging_lp.py— standardlp.Sersic+lp_linear.Sersicvariants, FD-validated (rel err ≤ 1e-5 / ≤ 3e-9). Evaluation point anchored near simulator truth so the positive-only NNLS keeps the source live; asserts the source block is live. Fixes the staleimaging/simpledataset pointer that broke the script on clean checkouts.scripts/jax_grad/imaging_mge.py— MGE source FD-validated through the linear inversion (rel err ≤ 5e-7). Fixes the staleimaging/source_complexpointer.scripts/jax_grad/imaging_pixelization.py(new) —RectangularUniform: strict FD assertions on all 14 params (AD = FD to 7 s.f.).RectangularAdaptDensity(os_pix=1): lens-light FD assertions + staircase-invariance assertions (LL bit-identical under ≤1e-6 mass shifts — autodiff's zero mass gradients are the correct a.e. derivative; fails loudly if mesh differentiability ever changes).scripts/jax_grad/point_source.py(new) —FitPositionsSourcesource-plane χ² FD-validated (rel err ≤ 5e-6, incl. magnification-via-Hessian); eager path (forward jit blocked upstream by theGrid2DIrregularxp gap).scripts/jax_grad/weak.py(new) —FitWeakFD-validated (rel err ≤ 3e-9), plain + per-galaxy redshift-scaled.Validation
All six scripts run green from the repo root on CPU JAX (float64), 2026-07-09, against current mains. Heart YELLOW acknowledged by the user at ship (pre-existing mcmc-smoke failures / worktree drift / assistant version skew — unrelated to this change).
🤖 Generated with Claude Code
https://claude.ai/code/session_01XGDp54jKUd3kUziTF77pNu