fix: JAX-trace Sérsic stellar-mass CSE deflection path (#499)#500
Merged
Conversation
The cored-steep-ellipsoid (CSE) decomposition used by Sérsic stellar-mass deflections was only partially ported to JAX, so a jitted likelihood over a Sérsic mass profile crashed with TracerArrayConversionError (surfaced by group/features/advanced/mass_stellar_dark/chaining.py on the release stack). This is an incomplete-JAX-port bug, not a jax-0.10.2 regression — isothermal/ NFW deflect analytically and never exercised the CSE path under jit. P1 (unblock): thread `xp` through the full CSE deflection path — radial_grid_from, decompose_convergence_via_cse, axis_ratio, and _decompose_convergence_via_cse_from (vectorised coefficient matrix + jnp.linalg.lstsq on the JAX path; scipy retained for NumPy). cse_settings_from is rewritten branch-free (xp.where) so a tracer sersic_index traces, with total_cses/sample_points frozen to a static 50/80 (conservative max) for static jit shapes and the per-index dex ranges preserved exactly. NFW/SersicGradient decompose signatures accept `xp` for compatibility. Validated: mass_stellar_dark/chaining.py completes (0 tracebacks); 432 test_autogalaxy mass-profile tests pass; jit round-trip nojit==jit to ~1e-11; numpy-vs-jax parity <=2e-5 for n in [1,4]. NumPy Sérsic-mass deflections shift ~1e-3 (finer CSE count). See PyAutoGalaxy#499. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
JAX-trace the Sérsic stellar-mass CSE deflection path so a jitted likelihood over a Sérsic mass profile no longer crashes with
TracerArrayConversionError. Surfaced bygroup/features/advanced/mass_stellar_dark/chaining.pyfailing fast (26.6s) on the release stack. Classification: real library bug (incomplete CSE→JAX port), not a jax-0.10.2 regression — isothermal and NFW deflect analytically (HK24) and never exercised the CSE path under jit, so this was the first release script to jit a CSE mass profile.P1 (unblock) only; P2 (parity hardening) tracked as a follow-up (see below).
API Changes
None to the public profile API — no symbol removed or renamed. Internal CSE helpers (
decompose_convergence_via_cse,cse_settings_from,_decompose_convergence_via_cse_from) gain an optionalxp=npkwarg. Behaviour change: NumPy Sérsic-mass deflections shift at the ~1e-3 level because the CSE count rose (30/40 → static 50/80, i.e. finer → closer to the integral truth). See full details below.Test Plan
test_autogalaxy/profiles/mass/→ 432 passed (NumPy path intact).mass_stellar_dark/chaining.py(release profile,PYAUTO_TEST_MODE=1, jax 0.10.2) → completes, 0 tracebacks.deflections(xp=jnp)nojit == jit to ~1e-11.ag.mp.Sersicn∈[1,4]./start_workspaceif so).Full API Changes (for automation & release notes)
Changed Signature
MassProfileCSE.decompose_convergence_via_cse(grid_radii, xp=np)— addedxp(abstract +Sersic,SersicGradient,NFWoverrides).MassProfileCSE._decompose_convergence_via_cse_from(..., xp=np)— addedxp; usesjnp.linalg.lstsqwhenxpisjax.numpy,scipy.linalg.lstsqfor NumPy.sersic.cse_settings_from(..., xp=np)— addedxp; standard (mass_to_light_gradient<=0.5) branch rewritten branch-free.Changed Behaviour
ag.mp.Sersic/Exponential/DevVaucouleurs) now trace underjax.jit.n(CSEtotal_cses/sample_pointsfrozen to static 50/80).test_sersicintegral-tolerance asserts (1e-4/1e-3) still pass.Migration
Validation checklist (--auto run — plan was not pre-approved)
Known P2 follow-up (not in this PR)
n<1NumPy-vs-JAX parity gap of 3e-3 (ill-conditioned CSEscipyvsjnplstsq).ag.mp.Sersicdeflections toautolens_workspace_test/scripts/profiles_jit.py(harness covers only Isothermal/PowerLaw/NFW — the gap that hid this bug).Generated by the PyAutoLabs agent workflow.