fix: make PointMass and SMBH JAX-compatible - #554
Merged
Conversation
PointMass.deflections_yx_2d_from routed through the legacy radial_grid_from / _cartesian_grid_via_radial_from helpers, which return an ArrayIrregular wrapper under xp=jnp on the irregular PSF-evaluation grids every imaging fit uses - jnp.multiply rejects it. Rewritten with raw xp ops (same maths, parity ~1e-18); potential_2d_from and convergence_2d_from hardened the same way (the convergence central-pixel argmin code was dead - its assignment was commented out). SMBH.__init__ used np.sqrt on the traced mass, raising TracerArrayConversionError; now tracer-safe pow. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This was referenced Aug 6, 2026
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
al.mp.PointMassandal.mp.SMBHfailed every JAX-mode fit (user report on 2026.8.4.1; reproduced on main):PointMass.deflections_yx_2d_fromrouted through the legacyradial_grid_from→_cartesian_grid_via_radial_fromhelpers, which return anArrayIrregularwrapper underxp=jnpon the irregular PSF-evaluation grids every imaging fit uses —jnp.multiplyrejects it (TypeError: multiply requires ndarray or scalar arguments, got ArrayIrregular).SMBH.__init__computednp.sqrt(mass_angular / np.pi)— withmassa free parameter the instance is built inside the jit trace, sonp.sqrtreceives a tracer (jax.errors.TracerArrayConversionError).Deflections are rewritten with raw
xpops (mirroring howIsothermalSphalready works under JAX),potential_2d_from/convergence_2d_fromare hardened the same way (the convergence central-pixelargmincode was dead — its assignment has been commented out for a long time), and the SMBH Einstein radius uses tracer-safe** 0.5.Fixes #553.
API Changes
None — internal bug fixes only. NumPy-path behaviour is numerically identical (deflections match the previous implementation and the analytic point-mass formula to ~1e-18). The behavioural change is that JAX-mode fits containing
PointMass/SMBHnow run instead of raising.See full details below.
Test Plan
test_autogalaxy/full suite: 1017 passedtest_autolens/point/: 115 passed against this branchAnalysisImaging.log_likelihood_functionwith stockPointMass(Einstein radius free) and stockSMBH(mass free): finite likelihoods — both configurations reproduced the user-reported errors on unfixed main as controlsjax.gradfinite for allPointMass/SMBHparameters (a non-finiteIsothermal.ell_compsgradient at exactly(0.0, 0.0)exists with and without the point mass — pre-existing, unrelated)θ_E²/r² · (y, x): max abs diff 8.7e-19Full API Changes (for automation & release notes)
Changed Behaviour
PointMass.deflections_yx_2d_from/potential_2d_from/convergence_2d_from— now tracer-safe underxp=jnp(previously raisedTypeError/TracerArrayConversionErrorin JAX-mode fits); NumPy results numerically unchanged.SMBH.__init__— Einstein radius computed with tracer-safe** 0.5instead ofnp.sqrt; accepts a tracedmass.Generated by the PyAutoLabs agent workflow.