From e462cf1acf05d38e71224fa53ffff4ff104d0d37 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 30 Jul 2026 20:00:40 +0100 Subject: [PATCH] docs: correct stale source-plane xp-blocker claims in jax_profiling (#657 phase 2 fixed it) --- jax_profiling/gradient/README.md | 2 +- .../gradient/point_source/source_plane.py | 53 +++++++++++-------- jax_profiling/jit/point_source/image_plane.py | 15 ++++-- .../jit/point_source/source_plane.py | 14 ++--- 4 files changed, 51 insertions(+), 33 deletions(-) diff --git a/jax_profiling/gradient/README.md b/jax_profiling/gradient/README.md index a062b4e..12858db 100644 --- a/jax_profiling/gradient/README.md +++ b/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 | diff --git a/jax_profiling/gradient/point_source/source_plane.py b/jax_profiling/gradient/point_source/source_plane.py index d8a6c15..7bcbfb9 100644 --- a/jax_profiling/gradient/point_source/source_plane.py +++ b/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 diff --git a/jax_profiling/jit/point_source/image_plane.py b/jax_profiling/jit/point_source/image_plane.py index ae830f5..706df0f 100644 --- a/jax_profiling/jit/point_source/image_plane.py +++ b/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 ------------------------------ diff --git a/jax_profiling/jit/point_source/source_plane.py b/jax_profiling/jit/point_source/source_plane.py index 4065d4c..e8d65f3 100644 --- a/jax_profiling/jit/point_source/source_plane.py +++ b/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")