[Jammy2211]

Summary

Fixes #767. Since 2026.9.27.2, AnalysisPoint.gradient_mode = "forward" makes gradient searches differentiate with jax.jacfwd, and for FitPositionsImagePairAll / FitPositionsImagePairAllSolved that forward gradient was NaN in every component, so every MultiStartAdam start was non-finite. Reverse mode was finite.

Root cause. The PointSolver pads its image positions with inf sentinel rows (remove_infinities=False). square_distance computes (d - inf) ** 2: the value is right (inf, so log_p = -inf), but the forward tangent is 2 * (d - inf) * 0 = NaN, which exp in the log-sum-exp carried into the log likelihood. Reverse mode happened to drop it.

Fix. Under JAX, all_permutations_log_likelihoods wraps the likelihoods in a jax.custom_jvp:

• primal = the unchanged per-position computation, so jitted log likelihoods are bit-identical (CPU fiducial 7.743201200876812, the autolens_profiling FIDUCIAL_SOLVED_LOG_L_BY_BACKEND["cpu"] pin, checked);
• tangent = the JVP of a new _log_likelihoods_padding_safe, which scores padded rows on a finite placeholder and masks them back to -inf (double-where), vectorized over model positions.

The NumPy path takes the same old code with no wrapper. Two simpler fixes were tried and rejected:

1. A per-position double-where still produced NaN once jitted. XLA's algebraic simplifier plus fusion leaked the non-finite discarded branch. It was finite without jit and NaN with jit, and --xla_disable_hlo_passes=algsimp removed the NaN.
2. A fully vectorized primal fixed the gradient but moved the fiducial by 1 ULP (…811).

API Changes

None — internal. all_permutations_log_likelihoods keeps its signature and values. New private helpers: _log_sum_exp, _log_likelihoods_from, _log_likelihoods_padding_safe.

Test Plan

• New test_pair_all.py::test__forward_mode_gradient_with_padded_model_positions__finite_and_matches_reverse_and_fd, for both classes. A smooth mocked solver with inf padding is jitted, and over PRNGKey(0..7) jacfwd must be finite, non-zero, equal to grad (rtol 1e-6) and equal to a central FD (rtol 1e-6). FD is checked on the mock only because the real solver's primal is discretised at pixel_scale_precision.
• New test__padded_model_positions__log_likelihoods_identical_to_padding_removed: with padding, the log likelihoods are identical to those without it (NumPy exact, JAX 1e-14).
• test_analysis_point_gradient_mode.py parity (subprocess, Fitness + register_model, 17 vectors = medians + PRNGKey(0..15)) is now parametrised over FitPositionsSourceSolved, FitPositionsImagePairAllSolved and FitPositionsImagePairAll with a real PointSolver. It asserts forward is finite and forward == reverse == Fitness.grad (rtol 1e-6 for the image-plane fits, where implicit-diff fwd/rev agree to ~1e-10, and 1e-8 for source-plane as before). It also asserts the gradient is non-zero.
• Red on unfixed source: with pair_all.py reverted, the 4 new image-plane tests fail and the source-plane control passes.
• Probe sweep, 4 parameter sets × 2 classes, jitted: values bit-identical to main, reverse gradients unchanged (≤3e-12 abs), forward == reverse (max rel 8e-11).
• Issue reproduction script (Fitness.grad, declared forward): [-41.62, 53.13, 31.79, 2.80, 1781.74], which equals reverse. It was NaN×5 before.
• pytest test_autolens/ on CPU fp64: 825 passed, 1 xfailed. The point subset gave 200 passed, including the static_lattice tie pins.
• Downstream smoke (worktree stack): autolens_workspace_test/scripts/point_source/jax_grad/gradient.py and jax_likelihood/image_plane.py both pass.

Downstream users of FitPositionsImagePairAll(Solved) (no edits needed — no API change)

• autolens_workspace_test: scripts/point_source/jax_grad/gradient.py (reverse-mode jax.value_and_grad, which is why it never caught this), jax_likelihood/{image_plane,point,solver_extent}.py, visualization/{visualization,visualization_jax,modeling_visualization_jit}.py
• autolens_profiling: scripts/point_source_image/likelihood_breakdown/{image_plane,gpu_bottleneck_map,solver_config_sweep,static_lattice_ab,vertex_dedup_ab,_point_solver_stage_map}.py, likelihood_runtime/{image_plane,image_plane_solved}.py. The fiducial_bit_exact gates pin 7.743201200876812 (CPU), which is unchanged here. The GPU pin (…806) was not re-measured.
• autolens_workspace: scripts/point_source/{fit,modeling,plot}.py, point_source/features/*, guides/point_source_pairing.py, cluster/{modeling,likelihood_function}.py, multi_dataset/features/imaging_and_point_source/modeling.py

Heart override

Heart is RED (release validation FAILED (stage integrate)). That reason has nothing to do with this change. The human granted the AUTONOMY.md "Human override for Heart RED (development only)" in-session (2026-10-04), covering push + PR-open only. This PR must not be merged until it gets a separate human /prm with green CI.

Patch release worth considering: the defect is live for every user of 2026.9.27.2+ who runs a gradient search on the image-plane all-pairs fits.

🤖 Generated with Claude Code

https://claude.ai/code/session_01S11WE9oj7Mvkfhc4EPBnyN