fix: investigate JAX gradient mismatch at parameter 9

[Jammy2211]

Overview

Investigate the gross autodiff-versus-finite-difference mismatch at parameter index 9 in scripts/jax_grad/imaging_pixelization.py. Establish whether JAX 0.10.2 changed the autodiff result relative to 0.9.2 before accepting any finite-difference boundary explanation or changing the test.

Parameter index 9 has been identified against the live model as galaxies.lens.mass.einstein_radius. The task begins as a workspace diagnostic and escalates to PyAutoArray/PyAutoGalaxy only if the version comparison establishes a source-level autodiff regression.

Plan

• Reproduce the complete imaging pixelization gradient audit under JAX 0.10.2 and an isolated JAX 0.9.2 overlay.
• Identify the exact failing mesh/configuration and compare parameter 9 across versions.
• If autodiff is stable, characterize the finite-difference surface around parameter 9 and determine whether a boundary or measure-thin branch flip explains the result.
• Compare the sibling interferometer gradient audit for a shared mechanism.
• Apply only the evidence-supported fix: a documented workspace exclusion for a proven FD artifact, or a library-source repair for a real autodiff regression.
• Re-run the affected gradient audits under both JAX versions.

Detailed implementation plan

Affected Repositories

• autolens_workspace_test (primary, phase 1)
• PyAutoArray (conditional: only if a library-source regression is proven)
• PyAutoGalaxy (conditional: only if a library-source regression is proven)

Branch Survey

Repository	Current Branch	Dirty?	Worktree claim
./autolens_workspace_test	main	clean	claimed by release-validation-tail-burndown on feature/release-validation-tail-burndown

Suggested branch: feature/jax-grad-param9-mismatch

Worktree root: ~/Code/PyAutoLabs-wt/jax-grad-param9-mismatch/ (not created while the existing claim remains)

Implementation Steps

1. Capture Python, JAX, JAXLIB and plugin versions, then run the complete scripts/jax_grad/imaging_pixelization.py under the current JAX 0.10.2 development environment.
2. Install JAX/JAXLIB 0.9.2 into disposable PyAutoMind/tmp/ storage and rerun the identical workspace commit and dependency stack through that overlay, leaving the development venv unchanged.
3. Record the exact failing variant and the full param_names mapping; compare autodiff and finite-difference values for index 9 (galaxies.lens.mass.einstein_radius) between versions.
4. If autodiff is version-stable, run a targeted finite-difference step sweep and evaluate eager/JIT likelihood values at theta ± h to distinguish a smooth derivative from a discontinuity or measure-thin solver branch flip.
5. Run the corresponding adaptive/kernel/uniform paths in scripts/jax_grad/interferometer.py and determine whether its release failure shares the same mechanism.
6. Outcome gate:
   • Proven FD artifact: add a narrowly documented skip_indices=(9,) only at the affected assertion, preserving the printed mismatch and explaining the boundary evidence.
   • Version-dependent autodiff regression: do not change skip_indices; identify the owning PyAutoArray/PyAutoGalaxy differentiation path, re-survey those repos, and implement a source fix with downstream workspace validation.
7. Re-run the affected gradient scripts under JAX 0.10.2 and 0.9.2, then run the relevant workspace smoke validation before shipping.

Key Files

• scripts/jax_grad/imaging_pixelization.py — failing imaging gradient audit and outcome-dependent workspace fix.
• scripts/jax_grad/interferometer.py — sibling failure comparison.
• scripts/jax_grad/util.py — finite-difference comparison and documented exclusion semantics.
• autoarray/inversion/ — conditional source investigation if JAX-version regression is proven.
• autogalaxy/ inversion/JAX likelihood path — conditional source investigation if ownership lies above AutoArray.

Original Prompt

Click to expand starting prompt

jax_grad imaging pixelization: autodiff vs finite-difference mismatch at parameter 9

Type: bug
Target: autolens_workspace_test
Repos:

• autolens_workspace_test
• PyAutoArray
• PyAutoGalaxy
  Difficulty: medium
  Autonomy: supervised
  Priority: normal
  Status: formalised

Surfaced while burning down the 2026-07-13 release-validation tail (PyAutoHeart#72). Once
the per-script timeout was raised (mode=release now allows these to complete — PyAutoHeart
commit c19c948), autolens_workspace_test/scripts/jax_grad/imaging_pixelization.py runs to
completion (~6:38 on CPU) and then fails its autodiff-vs-finite-difference gradient
assertion:

util.assert_gradients_match(comparison)   # scripts/jax_grad/util.py:171
AssertionError: Autodiff vs finite-difference mismatch at parameter indices [9]:
  ad=[149808.61606031], fd=[-5282762.48254538], abs_err=[5432571.0986057], tolerance=[5282.76258255]

The mismatch is gross — opposite sign, ~35× magnitude — not a tolerance drift. The earlier
configs in the same script all FD-match and PASS (RectangularUniform; RectangularAdaptDensity
staircase invariance; RectangularAdaptImage + reg.Adapt os_pix=4; RectangularAdaptDensity
os_pix=4). Only the final compare_gradients / assert_gradients_match (line ~452) trips, at
parameter index 9.

Leading hypothesis: fd=-5.28e6 is characteristic of a finite-difference blow-up —
parameter 9's likelihood has a near-discontinuity / boundary where central-FD is unreliable,
so AD (+1.5e5) is probably the correct value and FD is the artifact. If so the fix is to add
parameter 9 to the documented skip_indices in the script (the gradient-audit README in
autolens_workspace_developer/jax_profiling/gradient/ documents such FD-unreliable cases),
NOT a library change. BUT this MUST be confirmed it is not a jax-0.10.2 autodiff
regression first: the dev venv is now jax 0.10.2 (bumped 2026-07-14, see
[[reference_laptop_gpu_jax_setup]]); reproduce the same script on jax 0.9.2 and check whether
parameter 9 matched there. If AD changed between jax versions, this is a real gradient
regression in the pixelized-inversion custom_jvp path (PyAutoArray/PyAutoGalaxy), not a test
fix — do not paper over it with skip_indices ([[feedback_no_silent_guards]]).

Identify which parameter index 9 is (print param_names[9]) and which configuration's
comparison fails. Cross-ref [[project_jax_gradient_audit_shipped]] (the 2026-07 gradient
audit that these assertions encode). Sibling jax_grad/interferometer.py also failed in the
same release run — check whether it shares the same parameter/root cause.

Repro: ~/venv/PyAuto (now jax 0.10.2), from autolens_workspace_test/,
python scripts/jax_grad/imaging_pixelization.py (~7 min; it self-simulates). Release run:
PyAutoHeart workspace-validation 29279095224 (TestPyPI 2026.7.13.1.dev65601).
See [[project_release_2026_07_13_blocked_3bugs]].

[Jammy2211]

Progress Update — 2026-07-14

Summary

The two-version reproduction rules out a JAX 0.10.2 autodiff regression. Parameter index 9 is galaxies.lens.mass.einstein_radius; both JAX/JAXLIB 0.9.2 and 0.10.2 return the identical autodiff gradient +149808.61606031 for the failing RectangularKernelAdaptDensity (os_pix=1, bandwidth=0.1) configuration.

The failure is finite-difference contamination from a measure-thin solver branch. JAX 0.9.2 finds a clean 1e-8 FD step and the entire script passes; JAX 0.10.2 lands on contaminated samples at every configured step and reproduces the release failure exactly.

Diagnostic evidence

• Controlled inputs: identical workspace commit 5ba43278074c8de54b41a57ef9104e3c2f32d119, live PyAuto source checkouts, NumPy 2.2.6, CPU device.
• JAX 0.10.2 FD sweep for parameter 9: +5.433864e8, -5.417407e7, -5.282762e6 for relative steps 1e-8, 1e-7, 1e-6; all disagree with AD.
• JAX 0.9.2 FD sweep: +1.498086e5 at 1e-8 (relative error 4.48e-9, matching AD), while the larger steps hit branch flips.
• ULP probe under 0.10.2: at the exact theta+h float for the 1e-8 step, the likelihood jumps by +17.40585856; the immediately adjacent float one ULP lower follows the autodiff tangent to about 6e-9. At 1e-7, eager evaluation is clean while the compiled minus sample jumps by +17.40593075. At 1e-6, the minus sample jumps by +17.40657751 while a two-ULP neighbor is clean.
• These isolated likelihood jumps explain the gross opposite-sign FD values; they do not represent a smooth derivative.
• scripts/jax_grad/interferometer.py passes locally on JAX 0.10.2, but its kernel-mesh FD sweep shows the same poisoned-step family (~7.887e4 to 7.887e5) with a clean alternative step available for every parameter. The release failure is therefore consistent with the same unstable FD sampling mechanism, but does not reproduce locally.

Decision

The evidence supports the prompt's workspace-only outcome: document parameter 9 as FD-unreliable and pass skip_indices=(9,) at the strict assertion for the os_pix=1 kernel-density variant. Keep the parameter in the printed comparison and explain the two-version/ULP evidence at the call site. Do not change PyAutoArray/PyAutoGalaxy autodiff code.

Remaining

• Identify parameter index 9 and the failing configuration.
• Reproduce imaging_pixelization.py on JAX 0.10.2 and 0.9.2.
• Distinguish FD branch contamination from an autodiff regression.
• Compare the interferometer sibling.
• Apply the documented workspace exclusion and rerun validation.
• Review and ship the workspace change.

Blocker

autolens_workspace_test remains claimed by release-validation-tail-burndown at clean commit d1d3cfe (one commit ahead of main). This issue is registered in PyAutoMind as planned and must not create a competing worktree until that claim is released or transferred.

[Jammy2211]

Ship checkpoint — Heart RED

The workspace fix is implemented locally and validated by complete imaging_pixelization.py runs under CPU JAX/JAXLIB 0.10.2 and 0.9.2. Both pass; the later os_pix=4 kernel variants retain strict all-parameter FD assertions.

The ship gate is RED with these exact blocking reasons:

• PyAutoFit: 1 commit(s) behind origin
• release validation FAILED (stage integrate)

This issue is a narrow corrective fix for release validation FAILED (stage integrate). Under the corrective-PR exception, a human may explicitly authorize commit, push, and PR-open for that exact quoted reason despite the organism-wide RED verdict. The exception does not acknowledge the unrelated PyAutoFit drift and does not permit merge or release.

Authorization requested

Authorize the corrective-PR exception for the exact RED reason release validation FAILED (stage integrate) so this task may commit, push, and open its pending-release PR. Until then, no workspace commit, push, or PR will be created.

[Jammy2211]

Smoke Test Results — 2026-07-14

All 52 runnable smoke tests passed across 6 workspaces; 3 entries were skipped by their declared no_run.yaml rules.

Workspace	Passed	Failed	Skipped	Total
autofit_workspace	8	0	2	10
autogalaxy_workspace	8	0	0	8
autolens_workspace	9	0	0	9
autolens_workspace_test	15	0	1	16
euclid_strong_lens_modeling_pipeline	6	0	0	6
HowToLens	6	0	0	6

[Jammy2211]

Corrective-PR authorization — 2026-07-14

Contemporaneous human authorization was received in this session:

Authorize the corrective-PR exception for release validation FAILED (stage integrate)

• Exact authorized RED reason: release validation FAILED (stage integrate)
• Specific corrective issue: fix: investigate JAX gradient mismatch at parameter 9 #164, jax_grad parameter-9 autodiff-vs-FD mismatch.
• Causal mapping: the failed release integration reaches scripts/jax_grad/imaging_pixelization.py; parameter 9 is galaxies.lens.mass.einstein_radius; JAX 0.9.2 and 0.10.2 produce the same autodiff value, while JAX 0.10.2's three exact FD samples cross a measure-thin solver branch flip. The diff documents and excludes only that named comparison for the affected os_pix=1 variant while retaining the measured output and all other gradient, liveness, and FoM assertions.
• Evidence: full imaging script passed on JAX/JAXLIB 0.9.2 and 0.10.2; sibling interferometer script passed on 0.10.2; 52 runnable smoke checks passed across all six workspaces with 0 failures and 3 declared skips; review verdict CLEAN.
• Validation plan: open one pending-release corrective PR; a human may separately merge it; then build fresh post-merge wheels, rerun release-integration validation on main, and obtain a fresh Heart verdict.
• Scope limit: this authorizes commit, push, and one PR only. It does not authorize merge, issue close, release, or release rehearsal.
• Sibling RED reason remains unaddressed: PyAutoFit: 1 commit(s) behind origin.

[Jammy2211]

Corrective PR opened — 2026-07-14

• PR: test: document JAX finite-difference branch flip #165
• Commit: 540e093de2313e330b179a876399aab3650955fd
• Parameter 9: galaxies.lens.mass.einstein_radius
• Diagnosis: finite-difference branch flip at the RectangularKernelAdaptDensity (os_pix=1, bandwidth=0.1) solver boundary, not a JAX 0.10.2 autodiff regression; JAX 0.9.2 and 0.10.2 return identical autodiff values.
• Validation: imaging script passed on both JAX versions; interferometer sibling passed on 0.10.2; all-six-workspace smoke gate passed 52/52 runnable checks with 3 declared skips; review CLEAN.
• State: awaiting-merge. The corrective exception has ended at PR-open; merge, issue close, and release remain separate human-gated actions, and the sibling Heart RED reason PyAutoFit: 1 commit(s) behind origin remains.

[Jammy2211]

Corrective PR merged — 2026-07-14

PR

• test: document JAX finite-difference branch flip #165
• Merge commit: 6c4ff26ae3b0305d92d81f82b926ccb96373c578

Gate result

All four GitHub smoke jobs passed before the merge:

• Python 3.12 push + PR jobs: PASS
• Python 3.13 push + PR jobs: PASS

The reviewed head remained 540e093de2313e330b179a876399aab3650955fd, and the PR was MERGEABLE / CLEAN at merge time.

Remaining recovery work

Issue #164 remains open. The corrective-PR recovery sequence still requires fresh post-merge wheels, a release-integration rerun on main, and a new Heart verdict. Release and issue close were not authorized by the merge instruction and have not been performed.

[Jammy2211]

Corrective validation GREEN — closing

The fresh mode=release validation run on the current main HEADs passed end-to-end:

• Stage 2 TestPyPI rehearsal v2026.7.14.1.dev66001 — all 5 wheels built + uploaded.
• Stage 3 wheel-based real-search integration — 542 passed, 0 failed, 0 timeout, 87 skipped.
• Ingested → release_ready: true; Heart is now YELLOW (no RED reasons remain).

This fix is on the validated main HEADs (commit_shas matched), so its corrective validation is satisfied. Closing.