[Jammy2211]

Summary

Fixes a wrong JAX gradient through the border relocator when the lens is centred on the mask's symmetry centre. This is the case for every cell evaluated at the prior medians. Closes #671.

What goes wrong: ellipse_params_via_border_pca_from sizes the relocation ellipse with a = max(|x′|) and b = max(|y′|). A point-symmetric lens (Isothermal + external shear) centred on the symmetry centre of the mask traces the border into exact ± pairs, so both maxima are exact ties. JAX's max gradient shares itself among the entries equal to the result. Under jit, XLA can recompute the operand inside a fused graph with different rounding, so that equality test is unstable. The reverse-mode gradient of b came out 2× or 0.5× its true value, while eager mode was correct. Every pixelized-source likelihood with use_border_relocator=True inherits this.

Measured (jitted Fitness, jax_test imaging, RectangularRTUAdaptImage 39x39, MGE-60 lens light, exact prior medians):

parameter	before (autodiff / FD)	after (autodiff / FD)
mass.einstein_radius	25.46 / 24.18	24.183 / 24.183 (1.4e-7)
mass.ell_comps_0	54.25 / 55.89	55.891 / 55.891
mass.ell_comps_1	41.68 / 41.92	41.922 / 41.922

Stage bisection: the traced grid had identical tangents under eager, jit and FD. The relocated grid diverged only under jit reverse mode, and inside it only b (eager 0.5392, jit 1.0908, FD 0.5392).

Fix: a new helper max_at_argmax_from takes the element at the argmax instead of a max reduction, so the whole gradient lands on one point, identically under eager and jit. Values are bit-identical to before. At a genuine tie, any single element's derivative is a valid one-sided derivative of the maximum.

Not changed: the kernel-CDF per-axis min/max and std(axis=0).min() in interpolator/rectangular.py are the same kind of reduction. They showed no eager/jit divergence in axis-aligned and round configurations, which do produce exact ties, so they are left as they are.

About the original #671 symptom (|g| ≈ 2e9 on the HST fiducial): that part is not a library bug. The 141x141 grid at pixelization over-sampling 1 has a point exactly at the Isothermal centre, where the psi softening (a 1e-8 core) makes the true derivative about 1e8, which FD cannot resolve. The evaluation-point fix belongs to the profiling campaign (autolens_profiling#509). Diagnosis is on #671.

API Changes

Bug fix to JAX gradients only; numerical values are unchanged.

• Added autoarray.inversion.mesh.border_relocator.max_at_argmax_from(values, xp=np), a tie-robust maximum.
• ellipse_params_via_border_pca_from uses it for the semi-axes a and b. Same values, corrected jit gradients.
  See full details below.

Test Plan

• New test__max_at_argmax_from__tied_maximum_gradient_lands_on_one_element: one-hot gradient on a tied maximum, eager and jit. jnp.max fails it with a 0.5/0.5 split.
• Full test_autoarray suite: 2268 passed, 4 xfailed.
• End-to-end jitted repro: mass/shear gradients match FD (table above). Repro scripts in Science/scratch/rectangular_gradient_rtu_sparse_2026-11-02/.
• Profiling Witness on the HST fiducial once autolens_profiling#509 moves the fiducial evaluation point off the grid (the odd grid's symmetry centre is a grid point).

Full API Changes (for automation & release notes)

Added

• autoarray.inversion.mesh.border_relocator.max_at_argmax_from(values, xp=np): maximum of a 1D array taken as values[xp.argmax(values)], so its gradient is one-hot even at ties.

Changed

• ellipse_params_via_border_pca_from: semi-axes a, b computed via max_at_argmax_from. Values are identical. Reverse-mode gradients under jit are now correct at tied extents (previously 2× / 0.5× off).

Migration

• None.

Generated by the PyAutoLabs agent workflow.

🤖 Generated with Claude Code