fix: finite forward-mode gradient for FitPositionsImagePairAll(Solved)

autolens/point/fit/positions/image/pair_all.py

@@ -113,30 +113,64 @@ def all_permutations_log_likelihoods(self) -> np.ndarray:
         exact region gradient searches must traverse to find the basin. The shifted form is mathematically
         identical wherever the literal form is finite, and stays finite (the max term contributes exactly 0
         after the shift) at arbitrarily large mismatch.
+
+        Padded model positions -- the solver's `inf` sentinel rows, kept because `model_data` solves with
+        `remove_infinities=False` -- contribute `log_p = -inf`, i.e. nothing to the sum. Their value is right,
+        but their forward-mode tangent is not: `square_distance` gives `(d - inf) ** 2` a tangent of
+        `2 * (d - inf) * 0 = NaN`, which `jax.jacfwd` -- the default `gradient_mode="forward"` of
+        `AnalysisPoint` -- carried through `exp` into every gradient component (PyAutoLens#767). Under JAX the
+        likelihoods are therefore wrapped in a `jax.custom_jvp`: the primal is the unchanged per-position
+        computation below (so jitted log likelihoods stay bit-identical), and the tangent is the JVP of
+        `_log_likelihoods_padding_safe`, which scores padded rows on a finite placeholder and masks them out.
         """
 
         model_data = self.model_data.array
 
-        def log_sum_exp(log_ps):
-            # `initial` covers the zero-model-positions case (an empty `log_ps`), where a bare `max`
-            # raises: the -inf sentinel is clamped to 0 below, giving `log(sum of nothing) = -inf`,
-            # exactly the literal form's result, which `chi_squared`'s `has_image` fallback replaces.
-            max_log_p = self._xp.max(log_ps, initial=-np.inf)
-            # With no finite model position every log_p is -inf, and shifting by a -inf max would
-            # produce NaN (`-inf - -inf`) inside `exp` — including under `jax.grad`, where a NaN in
-            # the branch `chi_squared`'s `xp.where` discards still poisons the gradient. Clamp the
-            # shift to 0 so this case reduces to the literal form's `log(0) = -inf`, which the
-            # `has_image` fallback in `chi_squared` then replaces.
-            max_log_p = self._xp.where(
-                self._xp.isfinite(max_log_p), max_log_p, 0.0
-            )
-            return max_log_p + self._xp.log(
-                self._xp.sum(self._xp.exp(log_ps - max_log_p))
+        if self._xp is np:
+            return self._log_likelihoods_from(model_data)
+
+        import jax
+
+        @jax.custom_jvp
+        def log_likelihoods(model_data):
+            return self._log_likelihoods_from(model_data)
+
+        @log_likelihoods.defjvp
+        def log_likelihoods_jvp(primals, tangents):
+            (model_data,), (model_data_dot,) = primals, tangents
+            _, log_likelihoods_dot = jax.jvp(
+                self._log_likelihoods_padding_safe, (model_data,), (model_data_dot,)
             )
+            return self._log_likelihoods_from(model_data), log_likelihoods_dot
+
+        return log_likelihoods(model_data)
 
+    def _log_sum_exp(self, log_ps):
+        """
+        Max-shifted `log(sum(exp(log_ps)))` over the model positions (see `all_permutations_log_likelihoods`).
+        """
+        # `initial` covers the zero-model-positions case (an empty `log_ps`), where a bare `max`
+        # raises: the -inf sentinel is clamped to 0 below, giving `log(sum of nothing) = -inf`,
+        # exactly the literal form's result, which `chi_squared`'s `has_image` fallback replaces.
+        max_log_p = self._xp.max(log_ps, initial=-np.inf)
+        # With no finite model position every log_p is -inf, and shifting by a -inf max would
+        # produce NaN (`-inf - -inf`) inside `exp` — including under `jax.grad`, where a NaN in
+        # the branch `chi_squared`'s `xp.where` discards still poisons the gradient. Clamp the
+        # shift to 0 so this case reduces to the literal form's `log(0) = -inf`, which the
+        # `has_image` fallback in `chi_squared` then replaces.
+        max_log_p = self._xp.where(self._xp.isfinite(max_log_p), max_log_p, 0.0)
+        return max_log_p + self._xp.log(
+            self._xp.sum(self._xp.exp(log_ps - max_log_p))
+        )
+
+    def _log_likelihoods_from(self, model_data) -> np.ndarray:
+        """
+        The per-observed-position log likelihoods for the given model positions (`inf` padding rows included):
+        the primal of `all_permutations_log_likelihoods`.
+        """
         return self._xp.array(
             [
-                log_sum_exp(
+                self._log_sum_exp(
                     self._xp.array(
                         [
                             self.log_p(
@@ -152,6 +186,34 @@ def log_sum_exp(log_ps):
             ]
         )
 
+    def _log_likelihoods_padding_safe(self, model_data) -> np.ndarray:
+        """
+        `_log_likelihoods_from` with a forward-mode tangent that stays finite on padded model positions, used
+        only for the tangent of `all_permutations_log_likelihoods` under JAX.
+
+        Padded (`inf`) rows are scored on a finite placeholder and masked back to `log_p = -inf` (the
+        double-`where` pattern), so no `inf` reaches `square_distance`. Each observed position is scored
+        against all model positions in one vectorized `log_p` call: with one scalar `where` per model
+        position instead, XLA's CPU fusion of the jitted forward-mode graph still leaked a non-finite tangent
+        out of the discarded branch.
+        """
+        is_valid = self._xp.isfinite(model_data).all(axis=1)
+        # Transposed so `square_distance`'s `coord2[0]` / `coord2[1]` index all x / all y coordinates.
+        safe_model_data = self._xp.where(is_valid[:, None], model_data, 0.0).T
+
+        return self._xp.array(
+            [
+                self._log_sum_exp(
+                    self._xp.where(
+                        is_valid,
+                        self.log_p(data_position, safe_model_data, sigma),
+                        -np.inf,
+                    )
+                )
+                for data_position, sigma in zip(self.data, self.noise_map)
+            ]
+        )
+
     @property
     def chi_squared(self) -> float:
         """

test_autolens/point/fit/positions/image/test_pair_all.py

@@ -296,3 +296,134 @@ def test__model_positions_present__chi_squared_unchanged(data, noise_map):
     )
 
     assert fit.chi_squared == -2.0 * -4.40375330990644
+
+
+# Real model positions of the forward-mode gradient tests below, followed by solver-style `inf`
+# padding rows (the `PointSolver` pads to a fixed size and is called with `remove_infinities=False`).
+_GRAD_BASE_POSITIONS = np.array(
+    [(-1.03, -1.09), (0.35, 1.63), (1.57, 0.43), (1.25, 1.22)]
+    + [(np.inf, np.inf)] * 6
+)
+_GRAD_DATA = al.Grid2DIrregular([(-1.0, -1.1), (0.4, 1.6), (1.6, 0.45), (1.2, 1.25)])
+_GRAD_NOISE_MAP = al.ArrayIrregular([0.05, 0.05, 0.05, 0.05])
+
+
+class _PaddedSolver:
+    """Model positions that move smoothly with `theta` (scale, shift_0, shift_1), with `inf` padding
+    rows whose tangent is exactly 0 -- the shape `PointSolver.solve` returns under JAX."""
+
+    def __init__(self, theta):
+        self.theta = theta
+
+    def solve(self, tracer, source_plane_coordinate, xp=np, plane_redshift=None, remove_infinities=True):
+        import jax.numpy as jnp
+
+        base = jnp.asarray(_GRAD_BASE_POSITIONS)
+        is_padding = ~jnp.isfinite(base).all(axis=1)
+        moved = jnp.where(is_padding[:, None], 0.0, base) * self.theta[0] + self.theta[1:3]
+        return type("Positions", (), {"array": jnp.where(is_padding[:, None], jnp.inf, moved)})()
+
+
+@pytest.mark.parametrize(
+    "fit_cls, point",
+    [
+        (al.FitPositionsImagePairAll, al.ps.Point(centre=(0.07, 0.07))),
+        (al.FitPositionsImagePairAllSolved, al.ps.PointSolved()),
+    ],
+)
+def test__forward_mode_gradient_with_padded_model_positions__finite_and_matches_reverse_and_fd(
+    fit_cls, point
+):
+    """
+    Regression (PyAutoLens#767): `AnalysisPoint` differentiates with `jax.jacfwd` by default, and the
+    forward gradient of the all-pairs likelihood was NaN in every component. The padded (`inf`) model
+    positions gave `square_distance` a tangent of `2 * (d - inf) * 0 = NaN`, which `exp` carried into
+    the log likelihood; reverse mode happened to discard it.
+
+    Jitted, as gradient searches run it (XLA's fusion of the jitted graph is where a per-position
+    guard still leaked NaN), over eight parameter draws: `jacfwd` must be finite, non-zero, equal to
+    `grad`, and equal to a central finite difference of the (smooth, mocked-solver) likelihood.
+    """
+    jax = pytest.importorskip("jax")
+    import jax.numpy as jnp
+
+    jax.config.update("jax_enable_x64", True)
+
+    lens = al.Galaxy(
+        redshift=0.5,
+        mass=al.mp.Isothermal(centre=(0.0, 0.0), einstein_radius=1.6, ell_comps=(0.05, 0.0)),
+    )
+    tracer = al.Tracer(galaxies=[lens, al.Galaxy(redshift=1.0, point_0=point)])
+
+    def log_likelihood(theta):
+        return fit_cls(
+            name="point_0",
+            data=_GRAD_DATA,
+            noise_map=_GRAD_NOISE_MAP,
+            tracer=tracer,
+            solver=_PaddedSolver(theta),
+            xp=jnp,
+        ).log_likelihood
+
+    value = jax.jit(log_likelihood)
+    forward = jax.jit(jax.jacfwd(log_likelihood))
+    reverse = jax.jit(jax.grad(log_likelihood))
+
+    for key in range(8):
+        theta = jnp.array([1.0, 0.0, 0.0]) + jax.random.uniform(
+            jax.random.PRNGKey(key), (3,), minval=-0.02, maxval=0.02
+        )
+
+        grad_forward = np.asarray(forward(theta))
+        grad_reverse = np.asarray(reverse(theta))
+
+        h = 1.0e-6
+        grad_fd = np.array(
+            [
+                (float(value(theta.at[i].add(h))) - float(value(theta.at[i].add(-h))))
+                / (2.0 * h)
+                for i in range(3)
+            ]
+        )
+
+        assert np.all(np.isfinite(grad_forward)), (key, grad_forward)
+        assert np.any(grad_forward != 0.0)
+
+        scale = np.max(np.abs(grad_reverse))
+        np.testing.assert_allclose(grad_forward, grad_reverse, rtol=1e-6, atol=1e-12 * scale)
+        np.testing.assert_allclose(grad_forward, grad_fd, rtol=1e-6, atol=1e-6 * scale)
+
+
+def test__padded_model_positions__log_likelihoods_identical_to_padding_removed():
+    """
+    The padding guard changes tangents only: under JAX (and NumPy) the per-position log likelihoods
+    with `inf` padding rows are exactly those of the real rows alone.
+    """
+    jax = pytest.importorskip("jax")
+    import jax.numpy as jnp
+
+    jax.config.update("jax_enable_x64", True)
+
+    real = al.Grid2DIrregular(_GRAD_BASE_POSITIONS[:4])
+
+    def fit_for(solver, xp):
+        return al.FitPositionsImagePairAll(
+            name="point_0",
+            data=_GRAD_DATA,
+            noise_map=_GRAD_NOISE_MAP,
+            tracer=tracer,
+            solver=solver,
+            xp=xp,
+        )
+
+    expected = fit_for(al.mock.MockPointSolver(real), np).all_permutations_log_likelihoods()
+
+    padded_np = fit_for(
+        al.mock.MockPointSolver(al.Grid2DIrregular(_GRAD_BASE_POSITIONS)), np
+    ).all_permutations_log_likelihoods()
+    padded_jax = jax.jit(
+        lambda theta: fit_for(_PaddedSolver(theta), jnp).all_permutations_log_likelihoods()
+    )(jnp.array([1.0, 0.0, 0.0]))
+
+    np.testing.assert_array_equal(padded_np, expected)
+    np.testing.assert_allclose(np.asarray(padded_jax), expected, rtol=1e-14)