fix: PowerLawCore.convergence_2d_from drops xp (Isothermal not jit-traceable) - #633
Merged
Merged
Conversation
Dropping xp made Isothermal.axis_ratio use np.minimum on JAX tracers, so convergence_2d_from could not be traced under jax.jit with traced ell_comps (TracerArrayConversionError). Adds a jax-free regression test. Closes #632 Co-Authored-By: Claude Opus 5.5 <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
PowerLawCore.convergence_2d_fromcalledself.convergence_func(grid_radius=grid_eta)without forwardingxp, soIsothermal.axis_ratiofell back tonp.minimumon JAX tracers andconvergence_2d_from(..., xp=jnp)raisedTracerArrayConversionErrorunderjax.jitwith tracedell_comps. This affects everyPowerLawCoresubclass that uses the shared method (Isothermal, PowerLaw, IsothermalCore, ...). The fix forwardsxp=xp; it was the onlyconvergence_func(...)call site that dropped it.Closes #632
API Changes
None — internal. No public signatures change; the NumPy path is numerically unchanged.
Test Plan
test__convergence_2d_from__forwards_xp_to_convergence_func(spyxpinto the undecorated method; assertsconvergence_funcreceives it). Red on unfixed main (receivednumpy), green with the fix; jax is not imported by the test.pytest test_autogalaxy/— 1244 passed.jax.jitoverell_compsgives[0.55477504 0.67082039], identical to NumPy;jax.gradw.r.t. e1 = -0.298 (finite, non-zero). Unfixed source raisesTracerArrayConversionError.🤖 Generated with Claude Code