Repository navigation
fix: wrong JAX gradient through rectangular mesh (RTU all backends, Bilinear CPU sparse) #671
Description
Overview
On the imaging/rectangular/hst fiducial cell (39x39, no --standard-size), the JAX gradient through the rectangular source mesh is wrong by many orders of magnitude while log L is correct to 1e-10. RTU (kernel-CDF) is wrong on both CPU sparse and A100 dense, and the two backends disagree in sign. Bilinear (rank-CDF) is right on A100 dense but wrong on CPU sparse. Gradient-based samplers on these meshes get garbage gradients. Found in autolens_profiling#509 (likelihood-lifecycle 6c).
Plan
- Write a small CPU test that reproduces the bug on a mock rectangular cell. It compares
jax.value_and_gradwith a central-difference sweep, for RTU on dense and sparse, and Bilinear on sparse. That separates the two suspected bugs before any fix. - Work out which parameters carry the huge gradient components. Then step through the mesh code in order (CDF transform, knot inverse, interpolation weights, mapping/curvature) and check the gradient after each step.
- Fix each cause in PyAutoArray without changing any log L value: same values to ≤1e-10 on the existing tests.
- Check against the prompt's Witness on the real 39x39 HST fiducial cell, CPU sparse and A100 dense, by running
gradient_breakdown.pyread-only from the p6c autolens_profiling checkout (no autolens_profiling changes in this task). - Tier: judge. Merge mode: human
/prm.
Detailed implementation plan
Affected Repositories
- PyAutoArray (primary, only claimed repo)
- autolens_profiling: used read-only for the Witness via the p6c worktree (claimed by likelihood-lifecycle-p6c-gradient-walls). Any profiling-side change goes onto p6c or a follow-up.
Branch Survey
| Repository | Current Branch | Dirty? |
|---|---|---|
| ./array/PyAutoArray | main (c7da8b0) | clean |
Suggested branch: feature/rectangular-gradient-rtu-sparse
Worktree: ~/Code/PyAutoLabs-wt/rectangular-gradient-rtu-sparse/
Implementation Steps
- Repro tests (
test_autoarray/inversion/): gradient vs FD plateau (steps 1e-4..1e-7, seeded directions) for RTU dense, RTU sparse and Bilinear sparse; dense vs sparse gradient agreement. CPU, fp64. Cover both the fiducial-like and standard-size meshes, because the 74 standard-size control legs agree. - Suspect 1, RTU on all backends. It's in
autoarray/inversion/mesh/interpolator/rectangular.py.create_transformsbuildsinv_transform = reverse_interp(F_knots, knots). Where a knot sits in the empty tails, consecutiveF_knotsdiffer by about 1e-12. The gradient ofjnp.interpwith respect to itsxpargument (theF_knotstable) is roughly 1/ΔF there, which would give |g| ≈ 1e9. That inverse feedssource_plane_mesh_gridand the node centres, so it reaches the regularization even though the forward mapping is smooth. Other things to check: themin/max/std().min()normalisation and thexp.clipat saturation. Likely fix: a ΔF floor or a monotone reparametrisation so the knot inverse is smooth at the tails, or stop the gradient through a table that doesn't affect the value. - Suspect 2, CPU sparse only. Bilinear is right on dense and wrong on sparse. The suspect is the sparse curvature and mapped-data path:
inversion/imaging/sparse.py,sparse_operator.curvature_matrix_diag_from/curvature_matrix_off_diag_from, andinversion_imaging_util.mapped_reconstructed_operated_data_via_sparse_operator_from. Look for scatter, segment-sum or gather derivatives that drop or double-count the weight terms, and compare against dense at matched values. - Fix, re-run the value tests (log L unchanged), then run the Witness on the HST fiducial (CPU sparse locally or on RAL, A100 dense on RAL with a human go).
Key Files
autoarray/inversion/mesh/interpolator/rectangular.py: kernel and rank CDF transforms, knot inverse, bilinear weightsautoarray/inversion/mesh/mesh/rectangular_rtu_adapt_density.py/rectangular_rtu_adapt_image.pyautoarray/inversion/inversion/imaging/sparse.py,inversion_imaging_util.py: sparse curvature and mapped data
Evidence
~/Science/Science/scratch/delaunay_hilbert_2026-11-01/rtu_vs_bilinear/(gradient.json per leg) andresults/rectangular_hst_rtu_fiducial/probe.json+ README.md- RAL 413867–413870; /mnt/ral/jnightin/profiling/autolens_profiling_p2dg/output/p6_gradient_rtu/
Original Prompt
Click to expand starting prompt
Wrong JAX gradient through the rectangular source mesh: RTU (all backends) and Bilinear on the CPU sparse path
Type: bug
Target: PyAutoArray
Repos:
- PyAutoArray
- autolens_profiling
Themes: - jax-gradient
Difficulty: medium
Autonomy: supervised
Priority: high
Status: formalised
Consequence: judge
Witness: on imaging/rectangular/hst fiducial (39x39, no --standard-size) value_and_grad agrees with the central-difference plateau (steps 1e-4..1e-7, 3 seeded directions) to <= 1e-3 relative for RectangularRTU on CPU sparse and A100 dense and for Bilinear on CPU sparse, and CPU/A100 gradients agree
Review-minutes: 10
Unattended: never
Filed: 2026-11-02
Found during likelihood-lifecycle phase 6c (autolens_profiling#509), RTU vs Bilinear side run. Libraries at tag 2026.10.25.1. Human: "yes" (file it), 2026-11-02.
Symptom
Cell imaging/rectangular/hst, fiducial 39x39 (no --standard-size), campaign leg-11 flags, gradient_breakdown.py --gradient, mesh via --rect-mesh.
- RTU (kernel-CDF): the likelihood is smooth. Along the 3 seeded unit directions the central difference plateaus at about -2895 / 1237 / 2530 over steps 1e-4..1e-7. Autodiff gives 5.6e7 / -1.6e8 / 5.6e7 on A100 dense (RAL 413870) and -1.8e8 / 5.9e7 / 7.1e7 on CPU sparse (RAL 413868), with |g| ≈ 2e9. The two backends differ in sign, while log L agrees to 1e-10.
- Bilinear (rank-CDF): A100 dense gives |g| 3.3e4, consistent with the FD at h=1e-3 (413869). CPU sparse gives |g| 6.5e10 and an analytic directional derivative of -2.4e9 against an FD of -4386 (413867).
- Control: on the 74 standard-size v2 campaign legs (RAL 413791-413830), CPU and A100 |g| are identical on every cell, including imaging/rectangular at all 4 instruments.
The value is right and only the gradient is wrong, so the bug is in the derivative path: a near-singular term (sqrt / division / where guard at a degenerate point) in the RTU kernel CDF, or in the sparse mapping-matrix construction that the CPU path uses. Which one is open.
Evidence
~/Science/Science/scratch/delaunay_hilbert_2026-11-01/rtu_vs_bilinear/{rtu,bilinear}/results/imaging/rectangular/hst/{jax_cpu_sparse_fp64/hpc_ral_cpu,jax_gpu_dense_fp64/hpc_a100}/gradient.json~/Science/Science/scratch/delaunay_hilbert_2026-11-01/results/rectangular_hst_rtu_fiducial/probe.json(13-step FD sweep + log L slices, A100); write-up in that folder's README.md.- RAL: /mnt/ral/jnightin/profiling/autolens_profiling_p2dg/output/p6_gradient_rtu/.
Suggested first step
Reduce to a CPU-reproducible unit test: per-parameter gradient vs FD on the fiducial cell to find which parameters carry the huge components (mesh weights, lens mass, or lens light), then bisect the RTU / sparse mapper path with jax.debug / jax.grad on intermediate quantities.
Activity
Jammy2211 commentedon Nov 2, 2026
Diagnosis: the 1e9 gradient is the evaluation point, not the mesh
Reproduced locally on the real
imaging/rectangular/hstfiducial (RTU dense, MGE-60, 39x39, prior medians): |g| = 2.77e9 (RAL: 1.97e9). All of the error is inmass.centre_0andmass.centre_1: autodiff 1.8e9 / 2.1e9 against FD −944 / −380. Every other parameter agrees with FD.Mechanism:
psi_frominautogalaxy/profiles/mass/total/isothermal.pysoftens the core assqrt(q²x² + y² + 1e-16), a core of radius 1e-8. At that grid point, d(deflection)/d(centre) ≈ −1.5e8, checked directly.Control: with the centre prior means nudged by 1e-5″, |g| drops to 2.4e5 and the mass-centre components fall back to the FD scale (−1681 vs −1274, −2028 vs −1026; the remaining gap is that point now sitting about 1.4e-5″ from the core).
Re-scope (human, 2026-11-02): this issue now targets the remaining real error. With MGE-60 lens light, autodiff and FD disagree by 2–5% on the Einstein radius, ellipticity and shear. This also happens on the even-grid
jax_testdataset, where no grid point is near the centre, and dense and sparse agree exactly (so it is not the sparse path). The fiducial evaluation point is a profiling-side fix, noted on autolens_profiling#509.Also seen, not yet chased: with a non-linear Sersic lens light, the sparse path's log L differs greatly from dense (−52010 vs −933, with the warning "No blurring_image provided").
Repro + results:
Science/scratch/rectangular_gradient_rtu_sparse_2026-11-02/(repro.py; env varsREAL,MGE,MEDIANS,NUDGE).Jammy2211 commentedon Nov 2, 2026
Library PR created
#672 (pending-release)
Root cause of the 2–5% error:
ellipse_params_via_border_pca_fromtook the border-ellipse semi-axes asmax(|x′|)/max(|y′|). A lens centred on the mask's symmetry centre (as at the prior medians) traces the border into exact ± pairs, so both maxima tie. JAX's tie-sharingmaxgradient is unstable underjit: the reverse-mode gradient ofbcame out 2× or 0.5× its true value, while eager mode was correct. MGE was a red herring; the error only showed through the jittedFitnesspath.Fix: a new helper,
max_at_argmax_from, takes the element at the argmax. Values are unchanged, and the gradient is one-hot under eager andjit.Verification:
mass.einstein_radiusnow matches FD (24.183 vs 24.183, previously 25.46).test_autoarraysuite passes (2268 passed).Remaining:
Jammy2211 commentedon Nov 2, 2026
Shipped
max_at_argmax_frommakes the border-ellipse extents tie-robust, so reverse-mode gradients underjitare correct when the lens sits on the mask's symmetry centre.