Merge pull request #685 from PyAutoLabs/feature/point-solver-padded-row-grads

fix: sanitize padded rows in the PointSolver implicit-diff rule for reverse mode

Author: Jammy2211

autolens/point/solver/implicit_diff.py

@@ -34,7 +34,13 @@
 
 Padded rows (the ``inf`` sentinels of the fixed ``MAX_CONTAINING_SIZE`` output) are
 constants of the output shape; their tangent is forced to zero so they cannot inject
-NaNs into the batch.
+NaNs into the batch. Because reverse mode transposes the rule into row-summed
+cotangents, masking the *output* alone is not enough: the padded rows' solve inputs
+are also sanitized (Jacobian evaluated at a real solved image, ``a_mat`` replaced by
+the identity, ``rhs`` zeroed) so no NaN ever enters the linear algebra. Evaluating
+the Jacobian at the padded rows' former ``(0, 0)`` placeholder was itself a NaN
+source — profile centres typically sit at the origin, where deflection Jacobians
+are singular (cluster host halos; PyAutoLens#678 phase B).
 
 Known limitation — free cosmology parameters: ``Tracer`` is registered with
 ``cosmology`` as ``no_flatten`` aux, so a cosmology carrying traced parameters (a free
@@ -83,7 +89,15 @@ def implicit_tangents_from(jac_alpha, dalpha, dbeta, finite, xp):
     """
     identity = xp.eye(2)
     a_mat = identity[None, :, :] - jac_alpha
+    # Sanitize padded rows before the solve: a non-finite padded ``a_mat`` row
+    # survives the output masking in forward mode but not in reverse mode,
+    # where the transpose solves against ``a_mat`` row-by-row and sums the
+    # cotangents — one NaN padded row contaminates every parameter's gradient.
+    # Real (finite) rows are untouched, so legitimate near-critical divergence
+    # is still surfaced.
+    a_mat = xp.where(finite[:, None, None], a_mat, identity[None, :, :])
     rhs = dalpha + dbeta[None, :]
+    rhs = xp.where(finite[:, None], rhs, 0.0)
     dtheta = xp.linalg.solve(a_mat, rhs[..., None])[..., 0]
     return xp.where(finite[:, None], dtheta, 0.0)
 
@@ -171,7 +185,17 @@ def solve_padded_jvp(primals, tangents):
 
         theta = solve_padded(tracer, beta)
         finite = xp.isfinite(theta).all(axis=1)
-        theta_safe = xp.where(finite[:, None], theta, 0.0)
+        # Padded rows are anchored at the first real solved image rather than
+        # (0, 0): profile centres typically sit at the origin, where deflection
+        # Jacobians are singular (an NFW's jacfwd at its own centre is NaN),
+        # and although the padded rows' tangents are masked below, reverse mode
+        # transposes the rule into row-summed cotangents — one NaN row poisons
+        # the gradient of every parameter. At a real image the Jacobian is
+        # finite and the padded rows' contributions are exactly zeroed. With
+        # zero solved images the anchor row is itself non-finite and the
+        # gradient is NaN — the likelihood is already invalid there.
+        anchor = theta[xp.argmax(finite)]
+        theta_safe = xp.where(finite[:, None], theta, anchor[None, :])
 
         def deflections_single(position, tracer_):
             return deflections_from(position[None, :], tracer_)[0]

test_autolens/point/triangles/test_implicit_diff.py

@@ -37,6 +37,34 @@ def test_implicit_tangents_solve_the_linear_system():
             np.testing.assert_array_equal(dtheta[i], 0.0)
 
 
+def test_implicit_tangents_padded_nan_rows_never_reach_the_solve():
+    """
+    Padded rows carry whatever the Jacobian evaluated at their placeholder position
+    produced — NaN when that position sits on a profile centre. The rule must
+    sanitize those rows before the solve (identity ``a_mat``, zero ``rhs``): in
+    reverse mode the transpose solves against ``a_mat`` row-by-row and sums the
+    cotangents, so a NaN padded row would contaminate every parameter's gradient
+    even though the forward output masks it (#678 phase B, cluster cells).
+    """
+    jac_alpha = np.array(
+        [
+            [[0.5, 0.0], [0.0, 0.5]],
+            [[np.nan, np.nan], [np.nan, np.nan]],
+        ]
+    )
+    dalpha = np.array([[1.0, 2.0], [np.nan, np.nan]])
+    dbeta = np.array([0.1, -0.2])
+    finite = np.array([True, False])
+
+    dtheta = implicit_diff.implicit_tangents_from(
+        jac_alpha=jac_alpha, dalpha=dalpha, dbeta=dbeta, finite=finite, xp=np
+    )
+
+    assert np.isfinite(dtheta).all()
+    np.testing.assert_allclose(dtheta[0], (dalpha[0] + dbeta) / 0.5, rtol=1e-12)
+    np.testing.assert_array_equal(dtheta[1], 0.0)
+
+
 def test_implicit_tangents_near_critical_diverge_unclamped():
     # det(I - J) -> 0: the tangent must diverge with the true solve, never be clamped.
     eps = 1e-12