test: add JAX subhalo-redshift reproducer for issue #498#79
Merged
Conversation
Adds scripts/jax_likelihood_functions/imaging/subhalo.py, which mirrors lp.py but adds a `subhalo` galaxy to the model. Runs two scenarios: - Scenario A: fixed subhalo redshift (z=0.55) -> PASS under jax.jit + vmap - Scenario B: free `af.UniformPrior(0.2, 0.9)` redshift -> raises TracerBoolConversionError, as reported by @qiuhan96 The script exits 0 today (Scenario B failing is the expected outcome). Once the underlying tracer_util bug is fixed in PyAutoLens it will exit 1 on the same path, prompting the assertion polarity flip and turning the script into a regression test.
This was referenced May 8, 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
Adds an integration-test reproducer for PyAutoLens issue #498, reported by @qiuhan96.
The bug: setting
subhalo.redshift = af.UniformPrior(...)raisesjax.errors.TracerBoolConversionErrorunderjax.jit, becauseautolens/lens/tracer_util.plane_redshifts_fromcalls Python'ssorted(galaxies, key=lambda g: g.redshift)on a list containing one traced redshift scalar.This script reproduces the failure cleanly and exits 0 while the bug is open. Once the fix lands in PyAutoLens, Scenario B will start passing and the script will exit 1, surfacing in
run_all_scripts.shso the assertion polarity is flipped and the same file becomes the regression test.Scripts Changed
scripts/jax_likelihood_functions/imaging/subhalo.py— new. Reproduces issue #498. Two scenarios:redshift_subhalo=0.55): vmap + jit-wrappedanalysis.fit_from, expected PASS, asserts JIT path matches NumPy path withinrtol=1e-4.redshift_subhalo=af.UniformPrior(0.2, 0.9)): same calls, expected to raiseTracerBoolConversionErrorfromtracer_util.py:46. PrintsREPRODUCED: issue #498on the expected raise,UNEXPECTED: bug appears fixed(and exits 1) once the upstream fix lands.Test Plan
TracerBoolConversionErrorattracer_util.py:46(sorted(galaxies, ...)) as expected/smoke_test)🤖 Generated with Claude Code