research: JAX autodiff gradients audit across likelihoods

[Jammy2211]

Overview

Audit JAX autodiff support across the PyAutoLens likelihoods so gradient-based inference (NUTS/HMC via BlackJAX, gradient-aware nested sampling) stands on validated ground. Covers parametric light profiles (Sérsic, linear Sérsic, MGE), pixelized source reconstructions (Delaunay, rectangular), the point-source source-plane χ², and the weak-lensing likelihood. Deliverables are a per-likelihood gradient status report, autodiff-vs-finite-difference correctness tests (none exist today — current jax_grad/ scripts assert finiteness only), and documented root causes where differentiation breaks.

Plan

• Build a shared finite-difference comparison helper in autolens_workspace_test/scripts/jax_grad/ (vmap-honest evaluation so pure_callback const-folding can't fake results).
• Phase 1 — light profiles: re-run existing gradient probes on current mains; upgrade imaging_lp.py/imaging_mge.py from finiteness-only to FD-correctness; add a standard (non-linear) Sérsic case; document any breaks.
• Phase 2 — pixelized sources: re-confirm and write up why full Delaunay gradients are infeasible today (scipy pure_callback has no JVP rule); validate rectangular-mesh gradients vs FD, characterising the intentional stop_gradient approximation (anchor paper: https://arxiv.org/abs/2606.30620); check gradient flow through the NNLS solve.
• Phase 3 — validate point-source source-plane χ² gradients (they work under value_and_grad despite the forward-JIT blocker) and probe + validate the weak-lensing FitWeak likelihood (xp=jnp path exists, never gradient-probed).
• Wrap-up: audit README in jax_profiling/gradient/, follow-up tasks (e.g. Delaunay custom_jvp, Grid2DIrregular xp gap) filed as ideas, not fixed here.

Detailed implementation plan

Affected Repositories

• autolens_workspace_developer (primary — gradient probes + audit report)
• autolens_workspace_test (FD correctness tests)
• Read-only: PyAutoArray, PyAutoGalaxy, PyAutoLens, autolens_profiling (claimed by in-flight tasks; any library fix escalates as a follow-up task)

Branch Survey

Repository	Current Branch	Dirty?
./autolens_workspace_developer	main	1 untracked
./autolens_workspace_test	main	clean
./autolens_profiling	main	claimed: profiling-preopt-campaign
./PyAutoArray	main	claimed: nnls-solver-optimization

Suggested branch: feature/jax-autodiff-gradients-audit
Worktree: ~/Code/PyAutoLabs-wt/jax-autodiff-gradients-audit/

Implementation Steps

1. Shared FD helper autolens_workspace_test/scripts/jax_grad/util.py: central finite differences of the log-likelihood per model parameter (step scaled to parameter magnitude), compared to jax.grad with per-parameter rel/abs tolerance report; evaluate via vmapped fitness (fitness._vmap), never single-JIT on concrete arrays.
2. Phase 1: re-run autolens_workspace_developer/jax_profiling/gradient/imaging/mge.py and autolens_workspace_test/scripts/jax_grad/{imaging_lp,imaging_mge}.py on current mains; upgrade both jax_grad scripts to FD comparison; add lp.Sersic (standard, non-linear) coverage alongside lp_linear.Sersic and MGE; document tracing/differentiation breaks with file:line root cause.
3. Phase 2a (Delaunay): reproduce the 3 PASS / 8 ERROR result of jax_profiling/gradient/imaging/delaunay.py — root cause jax.pure_callback → scipy.spatial.Delaunay at PyAutoArray/autoarray/inversion/mesh/interpolator/delaunay.py:80, no JVP rule; write up why full gradients are infeasible (frozen triangulation = discrete combinatorial structure; a custom_jvp zero-rule unblocks the chain but leaves mesh-position derivatives undefined at cell boundaries).
4. Phase 2b (rectangular): read arXiv:2606.30620 first; validate existing gradients vs FD, explicitly characterising which parameter directions the lax.stop_gradient in create_transforms drops; add jax_grad/imaging_pixelization.py with tolerances informed by that analysis.
5. Phase 2c: gradient flow through the linear-inversion NNLS/fnnls solve — extend the existing solver-level guard jax_assertions/nnls.py, don't duplicate it.
6. Phase 3a: jax_grad/point_source.py — FD-validate FitPositionsSource.log_likelihood gradients including the magnification-via-Hessian term.
7. Phase 3b: new probe jax_profiling/gradient/weak/likelihood.py (FitWeak, xp=jnp) + jax_grad/weak.py FD test.
8. Wrap-up: per-likelihood status README (jax_profiling/gradient/README.md), follow-up ideas filed to the backlog, ship both workspace repos.

Key Files

• autolens_workspace_developer/jax_profiling/gradient/ — existing step-by-step value_and_grad probes (imaging mge/pixelization/delaunay, interferometer mge, point_source ×2); docstrings hold prior findings to re-validate.
• autolens_workspace_test/scripts/jax_grad/{imaging_lp,imaging_mge}.py — finiteness-only tests to upgrade.
• autolens_workspace_test/scripts/jax_assertions/nnls.py — solver-level finite-grad guard to extend.
• PyAutoLens/autolens/weak/fit.py — FitWeak, traceable with xp=jnp, pytree registered (read-only).
• PyAutoArray/autoarray/inversion/mesh/interpolator/{delaunay,rectangular}.py — the two gradient-killers/fixes under audit (read-only).

Verification

• Every new/upgraded jax_grad/ script runs green from the autolens_workspace_test root on CPU JAX; existing finiteness checks keep passing.
• FD agreement reported per parameter; intentional stop_gradient deviations documented with measured magnitude, not hidden by loose tolerances.
• Delaunay failure reproduced on current mains before being documented as current fact.

Original Prompt

Click to expand starting prompt

You're working in the PyAutoLens or PyAutoArray codebase, using the usual profiling workspace, profiling agent, and PyAutoBrain flow. First, investigate the state of JAX autodiff support for Sérsic light profiles, linear Sérsic light profiles, and Multi-Gaussian Expansion light profiles. Profile what's there, identify what breaks tracing or differentiation, and add or refine automated tests comparing autodiff against finite differences. Second, investigate gradients for pixelized source reconstructions, starting with the Delaunay mesh. We expect that full gradients may not be feasible there. Confirm why and document where it fails. Then move to the rectangular mesh, where we think gradients might be possible. Before starting that part, ask me for the relevant paper if you don't already have it. Third, validate gradients for the source plane chi-squared used for point sources, and for the weak lensing likelihood.

[Jammy2211]

Phase 1 complete — parametric light profiles all FD-validated ✅ (+ weak lensing early)

Built the shared finite-difference comparison helper (autolens_workspace_test/scripts/jax_grad/util.py: central differences per parameter, autodiff on the eager likelihood, FD sweep on a jitted one guarded by an eager-vs-jit consistency check) and upgraded the jax_grad scripts from finiteness-only to correctness:

Likelihood	Result
Imaging, lp.Sersic (standard)	✅ FD-matched, rel err ≤ 1e-5 over 22 params (most ≤ 1e-9)
Imaging, lp_linear.Sersic	✅ FD-matched, rel err ≤ 3e-9 over 20 params — gradient flows correctly through the NNLS solve, incl. source-shape params
Imaging, MGE (20×lp_linear.Gaussian)	✅ FD-matched, rel err ≤ 5e-7; step-by-step probe 9/9 PASS on current mains
Weak lensing (FitWeak, plain + redshift-scaled)	✅ FD-matched, rel err ≤ 3e-9 — gradients through the deflection-Hessian shear are correct (phase 3 item, landed early)

Findings along the way:

1. The old finiteness tests were broken on a clean checkout — imaging_lp.py/imaging_mge.py pointed at datasets that no longer exist (imaging/simple, imaging/source_complex) while their auto-simulate fallback writes imaging/jax_test. Re-pointed and guarded with al.util.dataset.should_simulate.
2. NNLS dead zones are real and subtle: at prior-median parameters the linear solver zeroes components (e.g. the lens bulge's median effective_radius ≈ 15" floods a 3" mask and absorbs the source arcs → source amplitude 0 → all source and mass gradients legitimately zero). The gradients are correct — but a naive test at such a point exercises nothing, and gradient-based samplers will see genuine plateaus. The scripts now anchor near the simulator truth and assert the source block is live.
3. Probe re-validation on current mains: Delaunay unchanged (3 PASS / 8 ERROR — pure_callback has no JVP rule, delaunay.py:126); point-source source-plane unchanged (4/4 PASS).
4. Phase 2 discovery: the rectangular-interpolator exploding-gradient fix (PyAutoArray#281) was never merged — closed unmerged 2026-06-01, apparently staled by the interpolator refactor. Yet the RectangularAdaptDensity probe on main shows finite, non-zero full-pipeline gradients (9/10 PASS). Either the refactor independently fixed the explosion, or the gradients are finite-but-wrong — exactly what the phase 2 FD comparison will decide.

Next: phase 2 — rectangular-mesh FD correctness (anchor paper arXiv:2606.30620 noted: create_transforms is an implementation of its ray-guided CDF transform) and the Delaunay infeasibility write-up.

[Jammy2211]

Phases 2 + 3 complete — the audit's headline finding is in 🔬

Phase 2a — Delaunay (re-confirmed, documented)

Unchanged on current mains: 3 PASS / 8 ERROR — everything from the mapping matrix onward raises Pure callbacks do not support JVP from _jax_delaunay_tables (PyAutoArray .../interpolator/delaunay.py:126). Write-up in jax_profiling/gradient/README.md covers the three layers: (1) mechanically, the callback needs a jax.custom_jvp zero-rule wrapper before anything can differentiate; (2) that gives frozen-triangulation gradients (the Delaunay analogue of dropped bin-boundary terms); (3) fundamentally, connectivity changes are discrete, so the frozen gradient is actually the correct a.e. derivative — the real question is sampler tolerance of the measure-zero non-smooth set.

Phase 2b — Rectangular mesh (headline finding)

The FD comparison caught something finiteness never could. FD step-size scans + two controls give the full picture:

config	einstein_radius: autodiff	FD h=1e-7	FD h=1e-6	FD h=1e-5
RectangularUniform, os_pix=1	−5.868261e5	−5.868261e5	−5.868261e5	−5.868261e5
AdaptDensity, os_pix=1	0.0	0.0 — LL bit-identical	0.0 — bit-identical	+1.7e4 (jump artifact)
AdaptDensity, os_pix=4	+7.33e4	+7.57e4	jumps	jumps

• RectangularUniform is perfect: autodiff = FD to 7 significant figures on all 14 parameters. Ready for gradient-based inference.
• RectangularAdaptDensity with pixelization over-sampling 1 destroys smooth mass information: the rank-space CDF transform (create_transforms) is exactly invariant under order-preserving deformations of the traced grid, so the likelihood is piecewise-constant in every mass/shear parameter (bit-identical LL under ≤1e-6 shifts). Autodiff's zero gradient is correct — the naive FD values are discontinuity artifacts from discrete rank re-orderings. This is a likelihood-design property, not an autodiff bug. Gradient-based mass inference is impossible in this configuration.
• Over-sampling > 1 restores a genuine smooth mass gradient (sub-pixel strain between interp queries and knots) and autodiff tracks it.
• Relation to the anchor paper (arXiv:2606.30620): Enzi et al. build the CDF transform from a smooth GP-weighted density rather than empirical point ranks — that's precisely what keeps their RTU formulation fully differentiable. Our empirical-rank variant is the non-differentiable corner of the same idea.
• History note: PyAutoArray#281 (the old exploding-gradient fix) is moot on the refactored code — nothing to re-land.

jax_grad/imaging_pixelization.py now encodes all of this: strict FD assertions for Uniform, lens-light FD assertions + staircase-invariance assertions for AdaptDensity (so any future change to mesh differentiability fails loudly).

Phase 3 — point-source χ² + weak lensing (both validated ✅)

• Source-plane positions likelihood (FitPositionsSource): FD-matched, rel err ≤ 5e-6 including the magnification-via-Hessian term (third derivatives of the potential). Confirmed the probe's finding that value_and_grad works even though forward jax.jit is still blocked by the Grid2DIrregular xp gap. flux/H0 legitimately zero in positions-only fits (FD agrees).
• Weak lensing (FitWeak): FD-matched, rel err ≤ 3e-9, both plain and per-galaxy redshift-scaled variants — already reported in the phase 1 comment.
• Interferometer MGE probe also re-confirmed 9/9 PASS (FD test for interferometer left as follow-up).

Next: wrap-up — ship the two workspace repos (audit README + 4 new / 2 upgraded FD test scripts), knowledge capture, and follow-up task proposals.

[Jammy2211]

Shipped 🚀

Both PRs are open with pending-release, Heart YELLOW acknowledged at ship (pre-existing reasons, unrelated to this change):

• autolens_workspace_test#157 — the FD gradient correctness suite (jax_grad/: util.py + 4 new scripts + 2 upgraded, all running green against current mains).
• autolens_workspace_developer#88 — the audit README (jax_profiling/gradient/README.md): status table, rectangular staircase verdict with measured FD step-size scans, Delaunay infeasibility write-up.

Knowledge capture and follow-ups also landed:

• PyAutoMemory methods_wiki/concepts/autodiff-implicit-diff.md (fills the long-standing placeholder).
• Five follow-up candidates appended to the ideas backlog: Delaunay custom_jvp zero-rule wrapper; differentiable adaptive mesh via smooth-density CDF (the Enzi et al. formulation); Grid2DIrregular xp gap fix; interferometer MGE FD test; NUTS/HMC trial on the now-certified likelihoods.

Audit verdict in one paragraph

Every smooth PyAutoLens likelihood is FD-certified under jax.grad: Sérsic (standard + linear through NNLS), MGE, point-source source-plane χ² (incl. Hessian magnifications), weak lensing (plain + redshift-scaled), and RectangularUniform pixelized sources. The two structural holes are Delaunay (hard pure_callback no-JVP error — unblockable with a custom_jvp zero rule) and the headline discovery: RectangularAdaptDensity at pixelization over-sampling 1 has a likelihood that is exactly piecewise-constant in every mass/shear parameter — autodiff's zero is the correct derivative, the discretisation itself destroys smooth mass information, and over-sampling > 1 (or a smooth-density transform) is required for gradient-based mass inference on adaptive meshes.

[Jammy2211]

Extension (user-requested): RectangularAdaptImage + production over-sampling ✅

Ran full 14-parameter FD sweeps at over_sample_size_pixelization=4 — how these meshes are actually used — including the complete production shape for RectangularAdaptImage (reg.Adapt(), al.AdaptImages, border relocator, positive-only solver; mixed precision off since FD needs float64):

config (os_pix=4)	mass/shear agreement (AD vs FD h=1e-7)	lens light
RectangularAdaptImage production	≤ ~1% (worst: einstein_radius 0.93%)	6 digits
RectangularAdaptDensity	≤ ~3.3% (worst: einstein_radius)	6 digits

All parameters live, both variants green in jax_grad/imaging_pixelization.py (now 4 variants: Uniform strict, AdaptDensity os1 staircase-invariance, AdaptImage-production os4, AdaptDensity os4 — with a liveness assertion on every parameter at os_pix=4). The residual ~1–3% on einstein_radius is FD micro-staircase contamination, not autodiff error: FD drifts with step size while AD is h-consistent.

Bottom line: gradient-based mass inference is certified in the production configuration. Only the os_pix=1 adaptive corner is unusable (the staircase finding). Pushed to both PRs (workspace_test#157 be8ce41, workspace_developer#88 6cb096a); merging next.

[Jammy2211]

Post-close extension: interferometer gradients (user-requested) — verdict complete

Q: was interferometer grad ok? Light profiles and rectangular (sparse linalg mode via inversion)?

• Light profiles: yes, certified. Standard lp.Sersic and lp_linear.Sersic source models FD-match through the TransformerDFT visibility transform at rel err ≤ ~1e-6 across all parameters (mass, shear, source shape, linear amplitudes through NNLS).
• Rectangular in sparse linalg mode: the staircase applies, with no escape hatch. On the exact production path (TransformerDFT + apply_sparse_operator(use_jax=True) + RectangularAdaptDensity + reg.Adapt + AdaptImages), every mass/shear autodiff gradient is exactly zero — correctly so: interferometer pixelization has no over-sampling, so the mesh's rank-transform queries always coincide with its knots, which is precisely the imaging os_pix=1 staircase corner. FD confirms (zero everywhere, one ~1e-7-scale rank-swap micro-jump on einstein_radius). Since interferometer models have no lens light, there are no usable gradients at all in this configuration — and unlike imaging, over-sampling isn't available as a mitigation.
• The working alternative: RectangularUniform on the same sparse path is FD-certified — all 7 mass/shear gradients live, strict agreement (≤ 2.4e-7). Interferometer gradient-based inference should use it until a smooth-density transform (the arXiv:2606.30620 formulation, already in the ideas backlog) makes the adaptive meshes differentiable.

PRs: PyAutoLabs/autolens_workspace_test#158 (test, jax_grad/interferometer.py, 3 variants all green) + #89 (README rows). This also retires the "interferometer MGE FD test" backlog bullet's LP/rectangular half; MGE-specific FD remains open.