Correct stale source-plane xp-blocker claims in jax_profiling (#657)

jax_profiling/gradient/README.md

@@ -32,7 +32,7 @@ finite and non-zero.
 | Imaging, Delaunay pixelization | **works — frozen-tables a.e.-exact gradient** (SHIPPED 2026-07-26: `stop_gradient` on the tables callback input) | **yes** (2026-07-26: median 9.6e-6, max 2.1e-3 over 14 params at documented rtol=1e-2; `jax_grad/delaunay.py`) | everything differentiable runs in-graph (visibility-walk point location, barycentric weights, dual areas, split points); the callback returns only int32 tables, piecewise-constant in the vertices, so freezing them under differentiation drops nothing — the exact a.e. derivative. Residual FD scatter on mass/shear = FD steps crossing triangle-flip events (measure-zero likelihood discontinuities). Batched caveat: callback is `vmap_method="sequential"` — KNN meshes remain the vmap-throughput option. See "Delaunay gradients: SHIPPED" section below |
 | Imaging, `KNearestNeighbor` (Wendland kNN) + `reg.ConstantSplit` / `reg.AdaptSplit` | **works** | **yes** (2026-07-26: strict FD-step-sweep, all 14 params, rel err ≤ 3.3e-8; `jax_grad/knn.py` variants A/B) | the JAX-native Delaunay-family mesh: Hilbert image mesh + edge zeroing, no scipy callback anywhere in the graph — gradients flow through traced query points AND traced mesh vertices. **Split-family regularization only**: `reg.Constant`/`ConstantZeroth`/`Adapt` need `MeshGeometryDelaunay.neighbors` (a direct scipy call on the traced mesh grid) and raise `TracerArrayConversionError` under `jax.grad` — pinned as a negative test in the script. Science caveat: Wendland kNN historically underperforms Delaunay (kernel knobs, caustic smearing — see PyAutoArray#317 background) |
 | Imaging, `KNNBarycentric` + `reg.ConstantSplit` | **works** (gradients only) | **yes** (2026-07-26: strict FD-step-sweep, all 14 params, rel err ≤ 4.1e-7; `jax_grad/knn.py` variant C) | 3-nearest barycentric weights; slightly noisier FD than Wendland (3-NN-set swaps move weights discontinuously — measure-zero jump sites). **Mesh failed its science gate as a Delaunay replacement** (PyAutoArray#317: 2.2% log-evidence drift, ~5% of vertices unreachable) — certified for gradient correctness, not for production science |
-| Point source, source-plane χ² (`FitPositionsSource`) | **works** (probe 4/4 PASS; forward `jax.jit` still blocked by the `Grid2DIrregular` xp gap) | **yes** (2026-07-09, rel err ≤ 5e-6; `jax_grad/point_source.py`) | includes magnification-via-Hessian term (3rd derivatives of the potential); flux/H0 legitimately zero in positions-only fits |
+| Point source, source-plane χ² (`FitPositionsSource`) | **works** (probe 4/4 PASS; forward `jax.jit` now also works -- xp gap fixed in PyAutoLens#657 phase 2 / PyAutoArray#414) | **yes** (2026-07-09, rel err ≤ 5e-6; `jax_grad/point_source.py`) | includes magnification-via-Hessian term (3rd derivatives of the potential); flux/H0 legitimately zero in positions-only fits |
 | Point source, image-plane (`FitPositionsImagePairAll`) | prior probe: **not differentiable** | n/a | `PointSolver` triangle-tiling forward solve uses `jnp.where` masking + integer neighbour lookups |
 | Weak lensing (`FitWeak`, `xp=jnp`) | **works** | **yes** (2026-07-09, rel err ≤ 3e-9, plain + redshift-scaled; `jax_grad/weak.py`) | gradients through the deflection-Hessian shear derivation are correct; no step-by-step probe needed — full pipeline validated first try |
 

jax_profiling/gradient/point_source/source_plane.py

@@ -30,21 +30,30 @@
    positions, plus ``noise_normalization``.
 7. ``log_likelihood = -0.5 * (chi_squared + noise_normalization)``.
 
-Surprising finding: full pipeline IS differentiable
----------------------------------------------------
+Full pipeline: forward JIT and gradient both succeed
+-----------------------------------------------------
 ``source_plane.py`` reports that the *forward* JIT of
-``AnalysisPoint(FitPositionsSource).log_likelihood_function`` is blocked
-by ``Grid2DIrregular.grid_2d_via_deflection_grid_from`` not propagating
-``xp``.  That blocker still applies under ``jax.jit``, but
-``jax.value_and_grad`` does not require lowering at the same boundary
-and the full-pipeline stage below succeeds with a finite, non-zero
-gradient.  Concretely: NUTS / HMC samplers driven through ``Fitness.call``
-should work for the source-plane likelihood today, even though the
-forward JIT path remains broken.
+``AnalysisPoint(FitPositionsSource).log_likelihood_function`` now JITs
+cleanly: the ``Grid2DIrregular.grid_2d_via_deflection_grid_from``
+xp-propagation bug it used to hit was fixed in PyAutoLens#657 phase 2
+(PyAutoArray#414, merged 2026-07-27). This script confirms the same
+pipeline is also differentiable via ``jax.value_and_grad``, so both the
+forward-JIT and gradient paths are usable for the source-plane
+likelihood. Concretely: NUTS / HMC samplers driven through
+``Fitness.call``, and gradient-based optimizers driven through a jitted
+forward pass, both work for the source-plane likelihood today.
+
+The one remaining JIT gap in the point-source stack is the
+whole-fit-object case (``jax.jit(analysis.fit_from)``): the returned
+``Fit`` pytree holds a ``PointSolver`` leaf that is not a registered JAX
+type, which JAX rejects at the output boundary. That gap does not affect
+``log_likelihood_function`` (used here and in ``source_plane.py``, which
+return a plain scalar) and is tracked in
+``PyAutoPrompt/autolens/fit_point_pytree.md``.
 
 That makes this probe load-bearing in a different way than the imaging
-probes -- it surfaces a usable gradient path that the forward JIT
-profiler conceals.
+probes -- it confirms both the forward-JIT and gradient paths are usable
+for the source-plane likelihood.
 """
 
 import numpy as np
@@ -243,9 +252,10 @@ def test_grad(label, func, params):
 # This is the JIT-able prefix from source_plane.py: build a tracer,
 # compute deflections at each observed image-plane position, and
 # subtract to get the source-plane positions. Stays inside raw arrays
-# the whole way -- no Grid2DIrregular result crosses the trace
-# boundary, so the xp-propagation blocker in
-# Grid2DIrregular.grid_2d_via_deflection_grid_from is avoided.
+# the whole way, mirroring the prefix profiled in source_plane.py (the
+# xp-propagation bug this used to route around,
+# Grid2DIrregular.grid_2d_via_deflection_grid_from, was fixed in
+# PyAutoLens#657 phase 2 / PyAutoArray#414).
 # ---------------------------------------------------------------------------
 
 
@@ -352,13 +362,12 @@ def step_positions_chi_squared(params):
 print("PART C -- FULL PIPELINE GRADIENT (via Fitness)")
 print("=" * 70)
 
-# source_plane.py reports a forward-JIT blocker for this same
-# pipeline (Grid2DIrregular.grid_2d_via_deflection_grid_from does not
-# propagate xp). value_and_grad does not require lowering at the same
-# boundary, so this stage is expected to PASS with a finite gradient
-# even while the forward JIT path is broken. A future status-flip
-# would indicate either the forward JIT was fixed (and we now hit a
-# different gradient path) or the gradient regressed.
+# source_plane.py confirms this same pipeline (via
+# log_likelihood_function) now JITs cleanly forward -- the
+# Grid2DIrregular.grid_2d_via_deflection_grid_from xp-propagation bug
+# was fixed in PyAutoLens#657 phase 2 (PyAutoArray#414). This stage is
+# expected to PASS with a finite gradient too, confirming both the
+# forward-JIT and gradient paths work for the source-plane likelihood.
 
 from autofit.non_linear.fitness import Fitness
 

jax_profiling/jit/point_source/image_plane.py

@@ -11,10 +11,17 @@
 loop), pairs each model image with the closest observed image, and computes
 a chi-squared in image-plane coordinates.
 
-Unlike the source-plane variant (see ``source_plane.py``), the full
-image-plane pipeline IS JIT-traceable end-to-end because ``PointSolver``
-threads ``xp=jnp`` through every step and ``FitPositionsImagePairAll``
-constructs its model-data via JAX-friendly operations.
+The full image-plane pipeline is JIT-traceable end-to-end because
+``PointSolver`` threads ``xp=jnp`` through every step and
+``FitPositionsImagePairAll`` constructs its model-data via JAX-friendly
+operations. The source-plane variant (see ``source_plane.py``) now JITs
+too: the ``Grid2DIrregular.grid_2d_via_deflection_grid_from``
+xp-propagation bug it used to hit was fixed in PyAutoLens#657 phase 2
+(PyAutoArray#414). The remaining JIT gap for point-source fits is the
+whole-fit-object case (``jax.jit(analysis.fit_from)``), which fails
+because the returned ``Fit`` pytree holds a ``PointSolver`` leaf that is
+not a registered JAX type -- tracked in
+``PyAutoPrompt/autolens/fit_point_pytree.md``.
 
 Pytree-native parameter inputs
 ------------------------------

jax_profiling/jit/point_source/source_plane.py

@@ -212,7 +212,8 @@ def jit_profile(func, label, *args, n_repeats=10):
 
 
 # ===================================================================
-# PART B — Full-pipeline JIT (expected to fail — see module docstring)
+# PART B — Full-pipeline JIT (JITs cleanly: the xp-propagation bug was
+# fixed in PyAutoLens#657 phase 2 / PyAutoArray#414, merged 2026-07-27)
 # ===================================================================
 
 print("\n" + "=" * 70)
@@ -260,11 +261,12 @@ def full_pipeline_from_params(params_tree):
 # PART C — JIT-able prefix: tracer ray-trace of observed positions
 # ===================================================================
 #
-# Even though the full pipeline is blocked, the dominant work in the
-# source-plane likelihood — ray-tracing the observed image positions to
-# the source plane via the tracer's deflection field — IS JIT-traceable
-# when the input/output stay as raw arrays.  We profile that prefix here
-# so the JIT-able portion of the source-plane path is still measured.
+# PART B already JITs the full pipeline (the xp-propagation bug fixed in
+# PyAutoLens#657 phase 2 / PyAutoArray#414 no longer blocks it). This
+# section isolates the dominant compute within that pipeline — ray-tracing
+# the observed image positions to the source plane via the tracer's
+# deflection field — as a standalone raw-array prefix, so it can be
+# profiled and vmap'd independently of the full likelihood.
 
 print("\n" + "=" * 70)
 print("JIT-ABLE PREFIX: ray-trace observed positions to source plane")