fix: tie-robust border-ellipse extents so jit gradients are correct

autoarray/inversion/mesh/border_relocator.py

@@ -266,12 +266,36 @@ def ellipse_params_via_border_pca_from(border_grid, xp=np, eps=1e-12):
     yprime = -s * dx + c * dy
 
     # Semi-axes from maximal extent
-    a = xp.max(xp.abs(xprime)) + eps
-    b = xp.max(xp.abs(yprime)) + eps
+    a = max_at_argmax_from(xp.abs(xprime), xp=xp) + eps
+    b = max_at_argmax_from(xp.abs(yprime), xp=xp) + eps
 
     return origin, a, b, phi
 
 
+def max_at_argmax_from(values, xp=np):
+    """
+    The maximum of a 1D array, taken as the element at its argmax rather than by a `max`
+    reduction, so that its gradient lands on exactly one element even when the maximum is
+    tied.
+
+    A lens centred on the mask's symmetry centre traces the border into exact +/- pairs, so
+    the border extents are tied maxima. JAX's `max` gradient shares itself among the entries
+    equal to the result, and under `jit` that equality test is unstable once XLA fuses the
+    operand into a larger graph and recomputes it with different rounding: the reverse-mode
+    gradient of the semi-minor axis came out 2x or 0.5x the true value (PyAutoArray#671).
+    Gathering one element gives the same gradient under eager and `jit`. At a tie any single
+    element's derivative is a valid one-sided derivative of the maximum.
+
+    Parameters
+    ----------
+    values
+        1D array whose maximum is returned.
+    xp
+        numpy-like module (np or jax.numpy).
+    """
+    return values[xp.argmax(values)]
+
+
 def relocated_grid_via_ellipse_border_from(grid, origin, a, b, phi, xp=np, eps=1e-12):
     """
     Rotated ellipse centered at origin with semi-axes a (major, x'), b (minor, y'),

test_autoarray/inversion/pixelization/test_border_relocator.py

@@ -5,6 +5,7 @@
 
 from autoarray.inversion.mesh.border_relocator import (
     ellipse_params_via_border_pca_from,
+    max_at_argmax_from,
     sub_border_pixel_slim_indexes_from,
 )
 
@@ -378,6 +379,7 @@ def test__relocated_grid_from__positive_origin_included_in_relocate():
 
     assert relocated_grid.over_sampled[1] == pytest.approx([1.95, 1.0], 1e-4)
 
+
 def test__ellipse_params__near_isotropic_border_uses_deterministic_axis():
     border_grid = np.array(
         [
@@ -414,3 +416,27 @@ def test__ellipse_params__anisotropic_border_retains_pca_major_axis():
     assert abs(float(phi)) == pytest.approx(np.pi / 2.0)
     assert float(a) > float(b)
 
+
+def test__max_at_argmax_from__tied_maximum_gradient_lands_on_one_element():
+    """
+    The border extents of a lens centred on the mask's symmetry centre are tied maxima.
+    JAX's `max` reduction shares the gradient among tied entries found by equality, which is
+    unstable under `jit` (PyAutoArray#671); the gradient must instead land on one element,
+    identically under eager and `jit`.
+    """
+    jax = pytest.importorskip("jax")
+    jax.config.update("jax_enable_x64", True)
+    import jax.numpy as jnp
+
+    values = jnp.array([0.3, 2.1, -0.4, 2.1, 1.7])
+
+    def maximum(values):
+        return max_at_argmax_from(values, xp=jnp)
+
+    assert float(maximum(values)) == pytest.approx(2.1)
+
+    for gradient in [jax.grad(maximum), jax.jit(jax.grad(maximum))]:
+        grad = np.asarray(gradient(values))
+
+        assert grad.sum() == pytest.approx(1.0)
+        assert np.count_nonzero(grad) == 1