Repository navigation
fix: jax.grad NaN at exactly-zero shear / multipole / ell_comps #631
Description
Overview
jax.grad returns NaN at exactly zero for ExternalShear (gamma_1, gamma_2), PowerLawMultipole.multipole_comps and ell_comps, so a gradient search starting at prior medians of 0 gets NaN gradients (the inference benchmark centres those priors on 1e-3 to dodge it). Cause: every polar conversion in autogalaxy/convert.py computes sqrt(c0**2 + c1**2), whose derivative is 0/0 at the origin.
Plan
- Add one private helper in
convert.pythat, on the JAX path only, nudges the x-component off the exact origin additively (gradient still flows, angle unchanged, forward value shifts ~1e-8) - Apply it at the top of every polar conversion (ell_comps, shear, multipole)
- Tests: finite jax.grad at (0,0) for ExternalShear / PowerLawMultipole / Isothermal, finite-difference match for the linear cases, fp32 finite check, NumPy values unchanged
- Run PyAutoGalaxy + PyAutoLens unit tests; ship library
- Follow-up (separate): drop the 1e-3 prior centring from the benchmark after release
Detailed implementation plan
Affected Repositories
- PyAutoGalaxy (primary)
Witness (main, fp64; jax.grad of a random-weighted deflection sum)
| profile | grad at (0,0) | grad at (1e-6,0) | central FD at 0 |
|---|---|---|---|
| ExternalShear | nan nan | 0.1816 0.4004 | 0.1816 0.4004 |
| PowerLawMultipole m=4 | nan nan | 0.5404 0.0910 | 0.5404 0.0910 |
| Isothermal ell_comps | nan nan | finite | — |
Branch Survey
| Repository | Current Branch | Dirty? |
|---|---|---|
| ./galaxy/PyAutoGalaxy | main | clean (3 behind origin) |
Claim: PyAutoGalaxy also claimed by workspace-config-cleanup (PyAutoMind#441) — its Galaxy PR #630 is merged, task awaiting release; parallel worktree human-approved 2026-09-27.
Suggested branch: feature/jax-grad-nan-zero-components
Implementation Steps
autogalaxy/convert.py: add_ORIGIN_NUDGE = 1e-8and_nudge_off_origin(c0, c1, xp)returningc1 + xp.where((c0 == 0) & (c1 == 0), _ORIGIN_NUDGE, 0.0)on the JAX path; NumPy path returnsc1unchanged (bitwise-identical values).- Apply it in
axis_ratio_and_angle_from,axis_ratio_from,angle_from,shear_magnitude_and_angle_from,shear_magnitude_from,shear_angle_from,multipole_k_m_and_phi_m_from(andmultipole_k_m_phi_m_and_shape_angle_fromif it does not route through it). Existing forwardwhereguards stay. - Tests (red-first on unfixed source): finite grad at (0,0) for the three profiles; grad == central FD for shear and multipole; fp32 finite; NumPy conversions still exactly (1.0, 0.0)/(0.0, 0.0) at the origin.
- PyAutoGalaxy + PyAutoLens pytest;
/ship_library.
Out of scope (recorded)
isothermal.py:99clampsaxis_ratio <= 0.99999, so d/d|ell_comps| is zero for |e| ≲ 5e-6 (finite, but no magnitude signal until the angle-direction step leaves zero).PowerLawMultipolem=1 singular at slope 2.- Benchmark 1e-3 prior centring removal in
autolens_inference— separate workspace task after release.
Key Files
autogalaxy/convert.py— polar conversionsautogalaxy/profiles/mass/sheets/external_shear.py,autogalaxy/profiles/mass/total/power_law_multipole.py— consumerstest_autogalaxy/test_convert.py— tests
Original Prompt
Click to expand starting prompt
jax.grad NaN at exactly zero components (ExternalShear, multipole_comps, ell_comps)
Type: bug
Target: @PyAutoGalaxy
Original request (verbatim)
fix this: - jax.grad returns NaN at exactly zero for ExternalShear, the multipole components and ell_comps. So a gradient search starting at prior medians of 0 gets NaN gradients. The benchmark centres those priors on 1e-3 to avoid it.
Witness (reproduced 2026-09-27 on PyAutoGalaxy main, fp64)
jax.grad of a random-weighted sum of deflections w.r.t. the (c0, c1) component pair:
- ExternalShear (0,0) -> [nan nan]; (1e-6,0) -> [0.1816 0.4004] == central FD at 0
- PowerLawMultipole m=4 (0,0) -> [nan nan]; (1e-6,0) -> [0.5404 0.0910] == central FD at 0
- Isothermal ell_comps (0,0) -> [nan nan]
Cause: sqrt(c0**2 + c1**2) in the autogalaxy/convert.py polar conversions has an
infinite/NaN derivative at the origin (the arctan2 already has a forward-value guard).
Origin: carried intake note from the point-source source-plane campaign (active.md).
Activity
Jammy2211 commentedon Sep 27, 2026
Library PRs created
convert.pyorigin nudge (traced JAX only,lax.select) plus removal of the Isothermal q ≤ 0.99999 clamp, replaced by a series form near q → 1Tests: galaxy 1258 passed; lens 762 passed, 1 xfailed. At (0,0), jax.grad now equals central finite differences for shear, the m=4 multipole and Isothermal ell_comps. Before the clamp removal, Isothermal was finite but wrong (−267 vs −0.53).
Workspace impact: none. The round-Isothermal parity scripts
tracer_jax.pyandprofiles_jit.pypass on the branch.Follow-ups filed in Mind:
Merge order: Galaxy first, then Lens (
/prm).Jammy2211 commentedon Sep 27, 2026
Shipped: merged PyAutoGalaxy #634 (c7fc595), then PyAutoLens #754 (897c2f76), with every CI leg green. Both are pending release. Follow-ups are filed in Mind: the dPIE / chameleon / MGE clamps and bare-sqrt sites (
draft/bug/autogalaxy/q_clamp_and_bare_sqrt_ellipticity_grad_sweep.md), and dropping the benchmark's 1e-3 prior centring after release (draft/feature/autolens_inference/drop_1e3_prior_centring_after_631.md).