fix: subhalo redshift as free parameter raises TracerBoolConversionError under JAX

[Jammy2211]

Overview

Reported on Slack by @qiuhan96 (working with an undergraduate at the University of Groningen) on the public pip release of autolens. Setting a subhalo's redshift as a free parameter (af.UniformPrior) raises jax.errors.TracerBoolConversionError during fitting. The only workaround today is use_jax=False, which is much slower. The cause is that autolens/lens/tracer_util.py performs Python-level <, <=, ==, float() on the redshift values, which are traced arrays under jax.jit when the subhalo redshift is free.

Plan

• Add a unit test that fits a model with a free-parameter subhalo redshift under JAX and asserts no TracerBoolConversionError.
• Reformulate grid_2d_at_redshift_from so the redshift comparisons used to choose where to insert the subhalo plane are JAX-friendly (e.g. compute candidate traced grids and select with jnp.where / jax.lax.switch, rather than branching on Python booleans).
• Stop calling float() and pairwise < on potentially-traced redshifts in plane_redshifts_from / planes_from — sort by the concrete redshifts of the non-subhalo galaxies (which are Python floats) and carry the subhalo redshift through as a traced scalar.
• Stop forcing the traced subhalo centre back through tuple(...) in AnalysisLens.tracer_via_instance_from (autolens/analysis/analysis/lens.py:116).
• Cover the edge cases in tests: subhalo before lens plane, between lens and source, equal to lens redshift, equal to source redshift.
• Verify the full Nautilus/Dynesty fit from the reporter's script runs end-to-end with use_jax=True.

Detailed implementation plan

Affected Repositories

• Jammy2211/PyAutoLens (primary)

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoLens	main	clean

Suggested branch: feature/subhalo-redshift-jax-fix

Implementation Steps

1. Reproduce. Add a regression test in test_autolens/lens/test_tracer_util.py (or a new test_autolens/analysis/test_analysis_lens_jax.py) that builds a 3-galaxy model (lens at z=0.5, subhalo with af.UniformPrior(0.2, 0.9), source at z=1.0) and calls jax.jit(analysis.fit_from)(instance), asserting it returns without raising.

2. Refactor autolens/lens/tracer_util.py:plane_redshifts_from. The list comprehension [float(galaxy.redshift) for galaxy in galaxies_ascending_redshift] (line 49) and the sorted(galaxies, key=lambda g: g.redshift) call (line 46) both fail on traced redshifts. Either:

   • Split inputs into "concrete-redshift galaxies" (sorted with Python sorted) and "traced-redshift galaxies" (the subhalo), and recombine without calling float(); or
   • Accept that this helper only ever sees Python-float redshifts and move the subhalo handling out before the call.

3. Refactor autolens/lens/tracer_util.py:grid_2d_at_redshift_from (line 199). The three branches that fail under JAX are:

   • Line 249: if redshift <= plane_redshifts[0]: return grid.copy()
   • Line 257: [... if galaxies[0].redshift == redshift]
   • Line 267-268: for plane_index, plane_redshift in enumerate(plane_redshifts): if redshift > plane_redshift: plane_index_insert = plane_index + 1

   Replace with a structured selection: compute the traced grid for every candidate insertion position (before plane 0, between each pair, after the last plane), then pick the right one via jax.lax.switch (or jnp.where over a stacked result) using a comparison vector built with jnp.less / jnp.less_equal. The numpy path keeps the existing implementation behind if xp is np:.

4. Update autolens/analysis/analysis/lens.py:99-116:AnalysisLens.tracer_via_instance_from. The line instance.galaxies.subhalo.mass.centre = tuple(subhalo_centre.in_list[0]) forces a Python tuple(...) of traced scalars. Either keep the centre as a traced 2-vector or skip the round-trip and pass the traced centre directly into the downstream Tracer build.

5. Tests. Cover four scenarios in test_autolens/lens/test_tracer_util.py:

   • subhalo redshift < lens.redshift
   • subhalo redshift == lens.redshift
   • lens.redshift < subhalo redshift < source.redshift (the typical case)
   • subhalo redshift == source.redshift

   Run each under both xp=np and xp=jnp (the latter inside jax.jit).

6. End-to-end check. Run the reporter's reproduction script (the model defined above plus a Nautilus/Dynesty search with a tiny nlive) with use_jax=True and confirm the fit completes.

Key Files

• autolens/lens/tracer_util.py — plane_redshifts_from, planes_from, grid_2d_at_redshift_from.
• autolens/analysis/analysis/lens.py — AnalysisLens.tracer_via_instance_from (subhalo branch at line 99-116).
• test_autolens/lens/test_tracer_util.py — new JAX regression tests.

Out of scope

• Changes to autogalaxy or autoarray.
• Anything related to the instance.perturb shortcut at lens.py:99 — that path is unaffected by the bug (subhalo redshift is taken from the model, not perturbed).

Original Prompt

Click to expand starting prompt

Free-parameter subhalo redshift breaks under JAX (TracerBoolConversionError)

Reporter

Reported on Slack by @qiuhan96 (with an undergraduate at the University of Groningen). Running the public pip release of autolens on Python 3.13.

Symptom

Setting the subhalo's redshift as a free parameter (a af.UniformPrior) raises a jax.errors.TracerBoolConversionError during model fitting. The fit only runs if use_jax=False is set, which is much slower.

Reproduction

import autofit as af
import autolens as al

bulge = al.model_util.mge_model_from(
    mask_radius=mask_radius,
    total_gaussians=20,
    gaussian_per_basis=2,
    centre_prior_is_uniform=True,
)

mass = af.Model(al.mp.Isothermal)
shear = af.Model(al.mp.ExternalShear)

lens = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass, shear=shear)

# Subhalo
subhalo_mass = af.Model(al.mp.IsothermalSph)
subhalo_mass.centre_0 = af.UniformPrior(lower_limit=-0.1, upper_limit=0.1)
subhalo_mass.centre_1 = af.UniformPrior(lower_limit=1.2, upper_limit=1.8)
subhalo_mass.einstein_radius = af.UniformPrior(lower_limit=0.01, upper_limit=0.4)

# Trigger: free-parameter redshift
redshift_subhalo = af.UniformPrior(lower_limit=0.2, upper_limit=0.9)
# redshift_subhalo = 0.6   # <-- works fine

subhalo_galaxy = af.Model(al.Galaxy, redshift=redshift_subhalo, mass=subhalo_mass)

# Source
bulge = al.model_util.mge_model_from(
    mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=False
)
source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)

model = af.Collection(galaxies=af.Collection(lens=lens, subhalo=subhalo_galaxy, source=source))

When redshift_subhalo is a UniformPrior, the fit raises:

File autolens/analysis/analysis/lens.py:99, in AnalysisLens.tracer_via_instance_from
    subhalo_centre = tracer_util.grid_2d_at_redshift_from(
        galaxies=instance.galaxies,
        redshift=instance.galaxies.subhalo.redshift,
        ...
    )
File autolens/lens/tracer_util.py:247, in grid_2d_at_redshift_from
    plane_redshifts = plane_redshifts_from(galaxies=galaxies)
File autolens/lens/tracer_util.py:46, in plane_redshifts_from
    galaxies_ascending_redshift = sorted(galaxies, key=lambda galaxy: galaxy.redshift)
...
jax._src.core.TracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[].

Workaround: setting use_jax=False lets the fit run, but is much slower.

Root cause

autolens/lens/tracer_util.py performs several Python-level operations on the redshift values that fail when one of them is a JAX traced scalar:

• Line 46 — sorted(galaxies, key=lambda g: g.redshift) does pairwise < comparisons on Python objects holding traced redshifts.
• Line 49 — [float(g.redshift) for g in ...] calls float() on a traced scalar.
• Line 249 — if redshift <= plane_redshifts[0]: is a Python branch on a traced boolean.
• Line 257 — [plane_index for ... if galaxies[0].redshift == redshift] filters on a traced boolean.
• Line 267-268 — for ...: if redshift > plane_redshift: plane_index_insert = plane_index + 1 again branches on a traced boolean and uses a Python integer to index the inserted plane.

grid_2d_at_redshift_from is called from autolens/analysis/analysis/lens.py:99 whenever instance.galaxies.subhalo exists, with redshift=instance.galaxies.subhalo.redshift. When that redshift is free, the value passed in is a traced array and every comparison above is illegal under jax.jit.

[Jammy2211]

Hey @qiuhan96 — picking this up now 👋

First step: a clean reproducer in our integration-test repo (autolens_workspace_test) that mirrors the model you shared but driven through our JAX likelihood-function harness, so we can see exactly what JAX is choking on and have a regression test for the fix to land against. I'll do a sanity-check run with the redshift fixed (your redshift_subhalo = 0.6 case) alongside the free-UniformPrior case so we can compare the two side-by-side.

Quick sidebar from the dig-in: I checked whether the SLaM subhalo helpers ever exposed a free_redshift-style API — they don't, the existing pipelines all pin subhalo.redshift to the lens redshift. So your af.UniformPrior(...) approach is the right one; the bug is on our side, not your model composition.

Will post back as soon as the reproducer is running. Thanks for the very thorough writeup — the traceback you included pointed straight at the broken line.

[Jammy2211]

Quick update @qiuhan96 — reproducer is up and we've got it pinned down 🎯

Built a script that mirrors your model but goes through our JAX integration-test harness, with two scenarios side-by-side:

• Scenario A — subhalo redshift fixed at z=0.55 → ✅ JAX vmap + JIT both work, log_likelihood ≈ -3.52e5, JIT matches NumPy within rtol=1e-4.
• Scenario B — subhalo redshift = af.UniformPrior(0.2, 0.9) → 💥 raises TracerBoolConversionError exactly as you reported.

The interesting bit: the failure is even earlier in the chain than the line your traceback flagged. It's not the if redshift <= plane_redshifts[0]: at tracer_util.py:249 — that line never gets reached. Python's built-in sorted() at line 46 already breaks because the key function returns traced scalars, and sorted does pairwise < to order them:

```
File autolens/lens/tracer_util.py:46, in plane_redshifts_from
galaxies_ascending_redshift = sorted(galaxies, key=lambda galaxy: galaxy.redshift)
jax.errors.TracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[].
```

So the fix has to handle the sort itself, not just the explicit <=/== branches further down. That actually simplifies things — we know the lens (z=0.5) and source (z=1.0) redshifts are concrete Python floats, and only the subhalo's redshift is traced, so we can sort the concrete-redshift galaxies once and merge in the traced subhalo with a JAX-friendly insertion that uses jnp.searchsorted / jax.lax.cond rather than a Python sort.

The reproducer is at autolens_workspace_test/scripts/jax_likelihood_functions/imaging/subhalo.py on branch feature/subhalo-redshift-jax-repro — once the fix lands, this script will turn into a regression test (Scenario B will be expected to PASS too).

In the meantime, your use_jax=False workaround is the right move. We'll let you know as soon as the fix is up. 🛠️

[Jammy2211]

Smoke Test Results — 2026-05-08

Workspace	Passed	Failed	Skipped	Total
autolens_workspace_test	11	0	1	12

All 11 smoke tests passed. 1 skipped (database/scrape/general — NEEDS_FIX 2026-04-27).

PR: PyAutoLabs/autolens_workspace_test#78

[Jammy2211]

Smoke Test Results — 2026-05-08

Workspace	Passed	Failed	Skipped	Total
autolens_workspace_test	11	0	1	12

All 11 smoke tests passed. (1 skipped: database/scrape/general.py — NEEDS_FIX 2026-04-27, pre-existing)

Tested against worktree feature/subhalo-redshift-jax-repro — new subhalo.py reproducer does not affect existing smoke suite.

[Jammy2211]

Reproducer PR is up: PyAutoLabs/autolens_workspace_test#79 ✅

Once that lands on main it'll be visible to anyone running our integration tests, and the same script becomes the regression check the moment the underlying tracer_util fix goes in (Scenario B will start passing, the script will exit 1, and we'll flip the assertion polarity in a follow-up commit).

Starting on the actual fix in PyAutoLens next — likely on a feature/subhalo-redshift-jax-fix branch. Will keep this thread updated.

[Jammy2211]

Reproducer is on main 🟢 — PyAutoLabs/autolens_workspace_test#79 merged.

So now we have a permanent, automated check that exercises this bug under our JAX integration-test runner. Anyone running run_all_scripts.sh will see it fail loudly the moment the upstream tracer_util fix accidentally regresses, and Scenario A's rtol=1e-4 jit-vs-numpy comparison guards against any drift in the multi-plane likelihood path on the way.

Now switching to the actual fix in PyAutoLens 🛠️

[Jammy2211]

Fix is in flight 🛠️

Plan we're going with:

• Make plane_redshifts_from, planes_from, and grid_2d_at_redshift_from JAX-aware. Each one keeps its existing numpy-fast-path when every galaxy redshift is a Python number (so concrete-redshift fits stay byte-for-byte identical), and gets a new JAX path when one or more redshifts are traced.
• The JAX path partitions galaxies into concrete vs traced, sorts the concrete ones with normal Python sort, and uses jnp.searchsorted(method='scan_unrolled') to find where the traced subhalo redshift slots into the concrete-plane sequence. grid_2d_at_redshift_from then runs the multi-plane ray trace at every candidate insertion position and selects the right one with jax.lax.switch.
• Also untangling the tuple(subhalo_centre.in_list[0]) line in analysis/lens.py:116 — that round-trip forces a traced 2-vector through a Python tuple(...), which would still break under JIT even after the tracer_util fix.

There's one open question we'll find out once the tracer_util block is past: whether the cosmology distance computations downstream are JAX-friendly with a traced redshift. They're called from traced_grid_2d_list_from and pull from astropy. If they're not, we'll file a follow-up — won't try to bundle that into this PR.

Library unit tests stay numpy-only; the JAX regression check stays in the reproducer script we just merged in autolens_workspace_test. Once the fix lands, we'll flip Scenario B's polarity from "expected to raise" to "expected to pass" in a small follow-up PR.

Will keep this thread updated as the PR goes up.

[Jammy2211]

PR is up: #499 ✅

Scenario B of our reproducer now passes — same log_likelihood = -3.523166e+05 as the fixed-redshift baseline, with the JIT path matching NumPy within rtol=1e-4. Existing JAX likelihood-function tests (lp.py, mge.py, etc.) all still produce their original regression literals, and the full PyAutoLens unit suite (273 tests) is green.

A few notes worth surfacing:

1. No public API changes. Every signature is unchanged. Existing scripts with concrete-float redshifts hit a fast-path that's byte-for-byte identical to the previous code, so no production fits change behaviour.

2. One small constraint added. When you make a Galaxy.redshift a free parameter under JAX, the galaxies must be declared in ascending-redshift order inside af.Collection(galaxies=...). Your bug-report code already does this (lens=, subhalo=, source=), and that's also the natural declaration order, so this should never bite anyone in practice. Documented in the PR body's "Migration" section.

3. Cosmology turned out to be fine. I'd flagged a risk that the astropy distance functions might also break on traced redshifts — they don't, at least for Planck15 (they accept traced scalars and return finite values). If a non-Planck cosmology turns out to need a JAX-friendly shim, we can file a follow-up, but it's not blocking you.

4. tuple(subhalo_centre.in_list[0]) in lens.py:116 works fine under JIT — turns out wrapping a (traced_y, traced_x) pair in tuple(...) is a no-op and the downstream xp.array(centre) builds a traced 1D array as expected. Initially flagged as suspect, but no fix needed there.

Once #499 is reviewed + merged we'll cut a release. After that, please retry your script with use_jax=True — it should run end-to-end. We'll also flip the polarity of the integration-test reproducer in autolens_workspace_test (small follow-up PR) so it serves as the regression test going forward.

[Jammy2211]

Smoke Test Results — 2026-05-08

Workspace	Passed	Skipped	Failed	Total
autolens_workspace_test	11	1	0	12

All 11 active smoke tests passed. database/scrape/general.py skipped (NEEDS_FIX per no_run.yaml).

This run was performed from worktree feature/multi-viz-imaging-small-datasets-override to verify the multi/visualization_imaging env_vars.yaml override fix before opening the PR.

[Jammy2211]

Smoke Test Results — 2026-05-08

Workspace	Passed	Failed	Skipped	Total
autolens_workspace_test	11	0	1	12

Run against PR #499's feature/subhalo-redshift-jax-fix branch with PYTHONPATH pointing at the patched library. Scripts covered: 7 JAX likelihood-function scripts (imaging/lp, mge, rectangular; interferometer/mge, rectangular; point_source/point; multi/mge), 3 aggregator scripts, and model_composition/multi_galaxy_mge.py. The single skip (database/scrape/general) is a pre-existing NEEDS_FIX unrelated to this change.

No regressions — every existing JAX likelihood-function script still produces its hardcoded regression literal within rtol=1e-4. The _any_traced guard at the top of each tracer_util function correctly routes concrete-redshift inputs to the unchanged numpy fast-path. ✅

[Jammy2211]

PR #499 is on main 🟢 — merge commit b790632.

Next step from our side: a tiny follow-up PR in autolens_workspace_test to flip the polarity of Scenario B in the reproducer script — right now it expects TracerBoolConversionError and exits 1 if Scenario B accidentally passes (i.e. now). Flipping it makes Scenario B's PASS the asserted outcome, turning the same script into the permanent JAX regression check for this bug. Will get that up shortly.

Once a release is cut from main and you've upgraded, please retry your model with use_jax=True and let us know how it goes. Thanks again for the great bug report — the traceback you pasted made this fix straightforward to scope. 🙏