fix: JAX-trace CSE mass-profile decomposition (Sersic/NFW deflections)

[Jammy2211]

Overview

group/features/advanced/mass_stellar_dark/chaining.py fails fast (~26 s) on the
release stack with a TracerArrayConversionError, not a timeout. Diagnosed as a
real library bug, not a jax-0.10.2 regression: the CSE decomposition path used
by Sérsic stellar-mass and NFW dark-mass deflections was only partially ported
to JAX (commit 35e6812), so it cannot be jax.jit-traced when the mass params are
free. It surfaced now because this is the first release script to jit a CSE mass
profile with free params (isothermal deflects analytically and never hits CSE).

Plan

• Thread xp into the deflection entry so the immediate crash clears.
• Make cse_settings_from traceable: xp ops + jnp.where/jnp.minimum instead
  of np.log10/np.min, and freeze total_cses/sample_points to a static
  conservative max (~50/80) so decomposition array shapes are constant under jit
  (no Python branch on the sersic_index tracer, no value-dependent shapes).
• Port _decompose_convergence_via_cse_from to xp.logspace, a vectorized
  coefficient-matrix build, and jnp.linalg.lstsq (replacing scipy.linalg.lstsq).
• Port both decompose_convergence_via_cse impls (Sérsic + NFW) to xp.
• Keep the NumPy path (xp=np) behaviourally identical; validate the JAX path via
  the autogalaxy_workspace_test parity scripts (jit round-trip + fitness._vmap),
  never test_autogalaxy/; confirm mass_stellar_dark/chaining.py runs green.

Phased: P1 — threading + traceable cse_settings_from/decompose to unblock the
release tail. P2 — parity hardening across the Sérsic-index range and NFW.

Detailed implementation plan

Affected Repositories

• PyAutoGalaxy (primary)
• PyAutoArray, PyAutoLens (reproduction only — no changes expected)

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoGalaxy	main	clean
./PyAutoArray	main	clean
./PyAutoLens	main	clean

Suggested branch: feature/mass-cse-jax-decompose

Root cause (reproduced on jax 0.10.2, JAX_TRACEBACK_FILTERING=off)

search[1] (lens light) cached; search[2] (lens mass = Sérsic + NFW) crashes on
its first jitted likelihood eval. The decorator AbstractMaker._xp
(PyAutoArray/.../decorators/abstract.py) resolves the backend from the xp passed
into the decorated call, so a dropped xp silently becomes numpy on a tracer.

1. cse.py:197 radial_grid_from(grid=grid, **kwargs) drops xp → np.sqrt on a
   float64[225] tracer.
2. After threading that, sersic.py:58 cse_settings_from → np.log10((23.0/ sersic_constant)**sersic_index) on a scalar tracer; whole decompose uses
   np.logspace on tracer bounds + scipy.linalg.lstsq (untraceable). NFW decompose
   (np.max(grid_radii), complex coord_func_g) has the same shape.

Implementation Steps (P1)

1. autogalaxy/profiles/mass/abstract/cse.py:197 — pass xp=xp to radial_grid_from.
2. autogalaxy/profiles/mass/stellar/sersic.py cse_settings_from (lines ~1–82) —
   remove Python branching on sersic_index; compute upper_dex/lower_dex with
   xp/jnp.where/jnp.minimum; return constant total_cses/sample_points
   (conservative max, ~50/80) so shapes are static.
3. autogalaxy/profiles/mass/abstract/cse.py _decompose_convergence_via_cse_from
   (lines ~85–147) — xp.logspace; build coefficient_matrix vectorized (no
   np.zeros+loop); jnp.linalg.lstsq when xp is jnp, else keep scipy for numpy
   parity/speed.
4. sersic.py:245 + dark/nfw.py:177 decompose_convergence_via_cse — xp the
   np.sqrt/np.max/np.real ops and radii bounds.

Validation

• autogalaxy_workspace_test JAX parity (jit round-trip + fitness._vmap) on a
  Sérsic mass + NFW model; NumPy path unchanged (test_autogalaxy/ stays green,
  NumPy-only).
• End-to-end: mass_stellar_dark/chaining.py under the release profile
  (PYAUTO_TEST_MODE=1) completes without a traceback.

Key Files

• autogalaxy/profiles/mass/abstract/cse.py — decomposition + deflection driver.
• autogalaxy/profiles/mass/stellar/sersic.py — cse_settings_from, Sérsic decompose.
• autogalaxy/profiles/mass/dark/nfw.py — NFW decompose.
• autogalaxy/profiles/geometry_profiles.py:113 — radial_grid_from (crash site).

Classification

Real library bug (incomplete JAX port), not a jax-0.10.2 regression. Surfaced
on release-validation run 29279095224 / TestPyPI 2026.7.13.1.dev65601
(PyAutoHeart#72).

Original Prompt

Click to expand starting prompt

group/mass_stellar_dark chaining fails with a JAX exception on release wheels

Type: bug
Target: autolens_workspace
Repos:

• autolens_workspace
• PyAutoGalaxy
• PyAutoArray
  Difficulty: hard
  Autonomy: supervised
  Priority: normal
  Status: diagnosed
  Fix-locus: PyAutoGalaxy (autogalaxy/profiles/mass — CSE decomposition path)
  Workflow: library
  Strategy: full jnp trace of the CSE decomposition; freeze total_cses/sample_points
  to a conservative static max (~50/80) on both paths (static shapes under jit);
  phased — P1 unblock chaining, P2 Sersic-range + NFW parity harden.

Surfaced burning down the 2026-07-13 release-validation tail (PyAutoHeart#72). On the
release wheels under the release profile:

scripts/group/features/advanced/mass_stellar_dark/chaining.py ... FAIL (26.6s)
  For simplicity, JAX has removed its internal frames from the traceback of the following
  exception. Set JAX_TRACEBACK_FILTERING=off to include these.

The run-runner captured only the filtered banner, not the actual exception — it is a real
JAX error during a group-scale mass_stellar_dark chaining pipeline (real sampler,
PYAUTO_TEST_MODE=1 for user workspaces), and fails fast (26.6 s), so it is NOT a timeout
(distinct from the other group//chaining scripts that TIMEOUT and are handled by the
release timeout bump PyAutoHeart#c19c948).

Reproduce from the autolens_workspace checkout on current main under the release profile,
with JAX_TRACEBACK_FILTERING=off to get the real traceback:

cd autolens_workspace
JAX_TRACEBACK_FILTERING=off PYAUTO_TEST_MODE=1 PYAUTO_SKIP_WORKSPACE_VERSION_CHECK=1 \
  python scripts/group/features/advanced/mass_stellar_dark/chaining.py

The dev venv is now jax 0.10.2 (matches the release stack — see
[[reference_laptop_gpu_jax_setup]]), so this should reproduce locally. Candidate causes given
the release-validation context: a jax-0.10.2 numerical/behaviour change in the chained
group-scale fit (cf. the jax-drift and LinAlg-not-PD siblings), a NaN/inf in a chained-prior
pass, or a real group/mass_stellar_dark modelling issue. Classify (jax regression vs real
bug) and fix at the right layer. Release run: PyAutoHeart workspace-validation 29279095224,
TestPyPI 2026.7.13.1.dev65601. See [[project_release_2026_07_13_blocked_3bugs]] and
PyAutoHeart#72.

Diagnosis (2026-07-14, reproduced locally on jax 0.10.2)

Classification: real library bug (incomplete JAX port of the CSE mass-profile
decomposition), NOT a jax-0.10.2 regression. jax 0.10.2 is incidental — the
release stack simply jits the likelihood; the same code fails on any jax version.

Reproduced on autolens_workspace main with JAX_TRACEBACK_FILTERING=off.
search[1] (lens light) is cached; search[2]__mass_stellar_dark crashes on the
first jitted likelihood eval. The lens mass is a Sersic (stellar) + NFW (dark)
model; the Sersic mass profile deflections go through the CSE decomposition
path, which is only partially JAX-threaded:

1. PyAutoArray .../decorators/abstract.py (AbstractMaker._xp) resolves the
   backend from the xp passed into the decorated call. So a dropped xp at a
   call site silently becomes xp=np.
2. First crash — PyAutoGalaxy/autogalaxy/profiles/mass/abstract/cse.py:197:
   self.radial_grid_from(grid=grid, **kwargs) drops xp (it is a named param
   of _deflections_2d_via_cse_from, not in **kwargs), so radial_grid_from
   (geometry_profiles.py:113 xp.sqrt(...)) runs with xp=np on a JAX tracer →
   TracerArrayConversionError (float64[225]). Threading xp=xp here clears it.
3. Second crash (after fixing Coordinates #2) — the decomposition itself is not traceable.
   With sersic_index free in search[2], decompose_convergence_via_cse runs on
   tracer params via numpy/scipy:
   • sersic.py:58 cse_settings_from → np.log10((23.0/sersic_constant)**sersic_index)
     on a scalar tracer → TracerArrayConversionError (float64[]).
   • sersic.py:280 np.sqrt(self.axis_ratio()); param-dependent radii_min/max.
   • cse.py:122/126 np.logspace(np.log10(radii_min), ...) on tracer bounds.
   • cse.py:143 scipy.linalg.lstsq — cannot trace at all.
     The dark NFW decompose (dark/nfw.py) has the same shape (np.max(grid_radii),
     scipy.lstsq) and would fail identically once reached.

Isothermal traces fine because it deflects analytically (xp.arctan/arctanh),
never touching CSE — so the CSE JAX port (commit 35e6812a, 2026-05-26 "port CSE
module to support JAX via xp parameter") was never actually exercised under jit
until this chaining script. This is the first release script to jit a CSE mass
profile with free params.

Fix locus: PyAutoGalaxy (profiles/mass/abstract/cse.py, stellar/sersic.py
incl. cse_settings_from, dark/nfw.py). Real fix ≈ thread xp through
cse_settings_from + both decompose_convergence_via_cse + _decompose_convergence_via_cse_from,
and replace scipy.linalg.lstsq/np.logspace with jnp.linalg.lstsq/xp.logspace
so the decomposition is jit-traceable. Validate via the JAX parity scripts in
autogalaxy_workspace_test (jit round-trip + fitness._vmap), never test_autogalaxy/.
Difficulty raised medium → hard (numerical path + lstsq-in-trace, affects all
CSE mass profiles under JAX).

[Jammy2211]

Ship sign-off requested (supervised --auto — ship parks per AUTONOMY.md)

P1 is implemented and validated in the worktree; changes are uncommitted pending your sign-off (supervised converts the ship-PR checkpoint into this question).

What was done (P1 — unblock)

JAX-traced the Sérsic stellar-mass CSE deflection path in PyAutoGalaxy. Diagnosis confirmed this is a real library bug (incomplete CSE→JAX port), not a jax-0.10.2 regression. Scope narrowed during implementation: NFW deflects analytically (HK24), not via CSE, so only the Sérsic path needed fixing for the deflection likelihood.

Files (feature/mass-cse-jax-decompose, 4 files, +104/−59):

• mass/abstract/cse.py — thread xp into radial_grid_from, decompose_convergence_via_cse, axis_ratio; port _decompose_convergence_via_cse_from (vectorised coefficient matrix, jnp.linalg.lstsq on the JAX path, scipy.linalg.lstsq retained for NumPy).
• mass/stellar/sersic.py — cse_settings_from rewritten branch-free (xp.where) so a tracer sersic_index traces; total_cses/sample_points frozen to a static 50/80 (conservative max) for static jit shapes; per-n dex ranges preserved exactly.
• mass/stellar/sersic_gradient.py, mass/dark/nfw.py — xp threaded so sibling signatures stay compatible (NFW convergence path also traceable as a bonus).

Validation (gate evidence)

• Tests: test_autogalaxy/profiles/mass/ → 432 passed (NumPy path intact).
• End-to-end: group/features/advanced/mass_stellar_dark/chaining.py under the release profile (PYAUTO_TEST_MODE=1, jax 0.10.2) → completes, 0 tracebacks (was FAIL 26.6s).
• JAX jit round-trip: nojit == jit to ~1e-11 (exact — the jit-correctness that was broken).
• NumPy-vs-JAX parity (x64): ≤2e-5 for Sérsic n∈[1,4].
• Review: self-review CLEAN — scoped, coherent, no unrelated changes.

API / behaviour

• No public symbol removed/renamed. Added optional xp=np kwarg to internal decompose_convergence_via_cse / cse_settings_from / _decompose_convergence_via_cse_from.
• Changed behaviour (NumPy): Sérsic-mass deflections shift at the ~1e-3 level for typical/high n because the CSE count rose (30/40 → 50/80, i.e. finer → closer to the integral truth; the test_sersic integral-tolerance asserts still pass at 1e-4/1e-3). No workspace script references these symbols directly, but downstream regression goldens that pin Sérsic-mass deflection values may need a ~1e-3 refresh (a /start_workspace follow-up if so).

Known P2 follow-up (agreed phasing — not in this PR)

• n<1 NumPy-vs-JAX parity gap of 3e-3 (ill-conditioned CSE scipy vs jnp lstsq).
• Add ag.mp.Sersic deflections to autolens_workspace_test/scripts/profiles_jit.py (parity harness currently covers Isothermal/PowerLaw/NFW only — the gap that hid this bug).

Heart state (evidence, not a branch gate for a feature ship)

Not GREEN: RELEASE VALIDATION integrate:fail (the ongoing 2026-07-13 tail burndown, organism-scope) + a stale 2026-07-09 TEST RUN 3f. Neither implicates this branch.

Decision requested

• Sign off → I commit, push, and open the pending-release PR (ends at PR-open; merge stays yours), then post the workspace-impact recommendation.
• Amend → tell me what to change first.
• Or fold the P2 parity-harden into this PR instead of a follow-up.

[Jammy2211]

Library PR created — #500 (pending-release)

Ship sign-off given in-session; commit 68a2a00d pushed on feature/mass-cse-jax-decompose, PR #500 open with pending-release. Run ends at PR-open (merge stays a human act).

Workspace impact: (iii) no API migration — cse_settings_from / decompose_convergence_via_cse are internal; no workspace script imports them. Caveat: the ~1e-3 NumPy behaviour change (finer CSE count) reaches three regression scripts that build Sérsic mass:

• autolens_workspace_test/scripts/mass_via_integral/sersic.py (CSE-vs-integral parity — most likely to notice)
• autolens_workspace_test/scripts/mass_via_integral/sersic_gradient.py
• autolens_workspace_test/scripts/mass/stellar.py

These assert against the integral ground truth (which the finer decomposition tracks at least as well), so they should still pass, but they want a smoke re-run; if any golden shifts ~1e-3, a small /start_workspace refresh PR follows.

P2 follow-up (separate task): (1) n<1 NumPy-vs-JAX parity gap 3e-3 (ill-conditioned CSE lstsq); (2) add ag.mp.Sersic deflections to profiles_jit.py (the parity-harness gap that hid this bug).

[Jammy2211]

Merged — PR #500 (squash b8bc0a3e), human-directed

The release-validation integration test that surfaced this is fixed. group/features/advanced/mass_stellar_dark/chaining.py (the FAIL in workspace-validation 29279095224) now runs green end-to-end (0 tracebacks, both searches complete) against the fix. Sérsic-mass parity regression scripts also pass at rtol=1e-3 (mass_via_integral/sersic.py, sersic_gradient.py).

P2 dropped (human call): the integration test is fixed, so we are not pursuing the n<1 lstsq parity-tightening or the profiles_jit.py Sérsic-harness addition as a follow-up. The n<1 NumPy-vs-JAX difference (3e-3) remains documented here for future reference but is below the CSE approximation's own error and is not release-gating.

Remaining to clear the release-validation RED (release-ops, human-gated): fresh post-merge wheels → rerun release-integration on main → new Heart verdict. Issue kept open until that confirms on wheels, matching the other 2026-07-13 tail fixes.

[Jammy2211]

Corrective validation GREEN — closing

The fresh mode=release validation run on the current main HEADs passed end-to-end:

• Stage 2 TestPyPI rehearsal v2026.7.14.1.dev66001 — all 5 wheels built + uploaded.
• Stage 3 wheel-based real-search integration — 542 passed, 0 failed, 0 timeout, 87 skipped.
• Ingested → release_ready: true; Heart is now YELLOW (no RED reasons remain).

This fix is on the validated main HEADs (commit_shas matched), so its corrective validation is satisfied. Closing.