From a79f27b910d9f4921a12f85fc0f756af5dd65c90 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Sat, 1 Aug 2026 13:20:24 +0100 Subject: [PATCH] fix: sanitize padded rows in the PointSolver implicit-diff rule for reverse mode MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Padded sentinel rows were zeroed to (0, 0) before the Jacobian evaluation — exactly where cluster profile centres sit, and an NFW's deflection Jacobian at its own centre is NaN. The forward tangent masks padded rows, but reverse mode transposes the rule into row-summed cotangents, so one NaN padded row poisoned every parameter's gradient (all-NaN grads in the #678 phase B cluster image_plane_solved cell; galaxy-scale cells were unaffected). Padded rows are now anchored at the first real solved image and the solve inputs are sanitized (identity a_mat, zero rhs) on non-finite rows only — legitimate near-critical divergence on real rows is still surfaced. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01DnTmLoJjJgMTze5uAbg1Jd --- autolens/point/solver/implicit_diff.py | 28 +++++++++++++++++-- .../point/triangles/test_implicit_diff.py | 28 +++++++++++++++++++ 2 files changed, 54 insertions(+), 2 deletions(-) diff --git a/autolens/point/solver/implicit_diff.py b/autolens/point/solver/implicit_diff.py index 263f01694..048a236de 100644 --- a/autolens/point/solver/implicit_diff.py +++ b/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] diff --git a/test_autolens/point/triangles/test_implicit_diff.py b/test_autolens/point/triangles/test_implicit_diff.py index c29718e74..f25d85ebb 100644 --- a/test_autolens/point/triangles/test_implicit_diff.py +++ b/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