Repository navigation
fix(latent): JAX jit assertion check drops every vis_lp latent #732
Description
Overview
The Euclid initial_lens_model/vis_lp stage (two-basis MGE lens light, SIE + shear, MGE source) of RAL job 342398 wrote no latent output on all ten euclid_dr1_prelim tiles: no files/latent/latent_summary.json, and search.log says compute_latent_samples: no finite latent samples remained after masking; skipping latent output. The vis_pix stage of the same tiles is fine.
The prompt's hypothesis (one latent NaN on every sample drops the whole block) is disproved by the engine as it ran on RAL: latent_samples_from already drops an all-NaN column and greedily salvages anti-correlated NaNs, so the log line can only mean every latent was NaN on every sample. Reproduced on the pipeline's committed simulated dataset with use_jax=True: LatentLens.variables / LatentEuclid.variables build the instance with model.instance_from_vector(vector=parameters), whose default assertion check applies a Python not to a traced boolean inside the per-sample jax.jit the latent batch path uses. The vis_lp model gained an ordering assertion on its two MGE bases on 2026-09-08 (order_bases=True, PyAutoFit#1583), so the trace raises TracerBoolConversionError on every sample; the engine's _safe_jitted swallows that into a NaN row, and the global mask then finds nothing. Fitness avoids exactly this on its JAX path with ignore_assertions=True, xp=jnp. A second, independent trace failure lives in the pipeline's own _source_flux_latents_on_uniform_grid (a Grid2D.from_mask(..., xp=jnp) inside the jit); with both fixed all twelve latents come out finite under jit and equal to the eager values.
Plan
- PyAutoLens: build the latent instance the way
Fitnessdoes on the JAX path (assertions skipped,xpthreaded); NumPy path unchanged (assertion violation still raisesFitException, dropped as a NaN row). - PyAutoFit: stop the engine hiding the raise. Count the per-sample failures, log the count and the first traceback, and when every sample raised say so instead of the generic "no finite latent samples remained".
- Tests: PyAutoLens unit tests for the helper on both backends plus an end-to-end
LatentLens.variablesunderjax.jitwith a model assertion; PyAutoFit tests asserting the warnings carry the cause. - Euclid pipeline follow-up (own prompt/PR):
LatentEuclid.variablesuses the new helper and the uniform source-flux grid is built on the static NumPy mask. Then avis_lprerun on tile 102005065 is the witness.
Detailed implementation plan
Affected Repositories
- PyAutoLens (primary)
- PyAutoFit
- euclid_strong_lens_modeling_pipeline (follow-up prompt, not this task)
Branch Survey
| Repository | Current Branch | Dirty? |
|---|---|---|
| ./PyAutoLens | main (7d1b04d, the commit RAL ran) | clean |
| ./PyAutoFit | main (6ea1dda; latent.py identical to RAL's 66f9f8d) | clean |
Suggested branch: feature/vis-lp-latent-jit-assertions (both repos)
Implementation Steps
PyAutoLens/autolens/analysis/latent.py: addlatent_instance_from(model, parameters, xp)—xp is np→model.instance_from_vector(vector=parameters)(assertions checked, unchanged); otherwisemodel.instance_from_vector(vector=parameters, ignore_assertions=True, xp=xp).LatentLens.variablescalls it; the class docstring tells subclasses that build their own instance to use it.PyAutoFit/autofit/non_linear/analysis/latent.py: inlatent_samples_from, afailurescounter + firsttraceback.format_exc()recorded by_safe_jitted/_safe_compute; after the batch loop warn "raised on N of M samples" with the traceback; the two "no finite latent samples remained" warnings route through_warn_no_latents(), which names the every-sample case ("the latent function is broken for this model").PyAutoLens/test_autolens/analysis/test_latent.py:test_latent_instance_from_numpy_checks_assertions,test_latent_instance_from_jax_skips_assertions_under_jit(asserts the default path raisesTracerBoolConversionErrorand the helper traces),test_latent_lens_variables_traces_under_jit_with_model_assertion(7x7 fixture,use_jax=True, values finite and equal to the NumPy evaluation). Both JAX testsimportorskip("jax").PyAutoFit/test_autofit/analysis/test_latent.py: caplog assertions ontest_new_path_skips_arbitrary_exception_samples; newtest_every_sample_raising_names_the_cause_instead_of_a_blank_block.
Key Files
PyAutoLens/autolens/analysis/latent.py—latent_instance_from,LatentLens.variablesPyAutoFit/autofit/non_linear/analysis/latent.py—latent_samples_fromfailure reportingPyAutoFit/autofit/non_linear/fitness.py:415-440— the JAX-path precedent being mirroredeuclid_strong_lens_modeling_pipeline/util.py:516,621— the two pipeline sites (follow-up)
Work Classification
Library
Worktree root
~/Code/PyAutoLabs-wt/vis-lp-latent-jit-assertions/ (this session is web-github: session clones, no worktree)
Original Prompt
Click to expand starting prompt
MGE lens-light stage writes no latent output: one non-finite latent drops all twelve
Type: bug
Target: PyAutoLens
Repos:
- PyAutoLens
Difficulty: medium
Autonomy: supervised
Priority: high
Status: formalised
Consequence: judge
Witness: thevis_lpsearch ofeuclid_dr1_prelimtile 102005065 (342398 dataset and config) writesfiles/latent/latent_summary.jsonwith every finite latent populated, and a unit test shows a latent that is NaN on every sample no longer suppresses the other latents.
Review-minutes: 20
Unattended: ready
Filed: 2026-09-10
Symptom
In PyAutoLens latent output, the MGE lens-light stage (initial_lens_model vis_lp in the euclid
pipeline) writes no latent output at all. All ten dr1_prelim initial_lens_model/vis_lp results
of RAL job 342398 have no files/latent/latent_summary.json and no latent.results; search.log
says
compute_latent_samples: no finite latent samples remained after masking; skipping latent output.
latent_samples_from masks globally, so one latent that is non-finite on every sample drops all
twelve latents for the search; the vis_pix stage of the same tiles is fine (12 keys, all
populated).
Do
- Diagnose which latent is non-finite for the two-basis MGE model and why (aperture-flux latents
on aBasisof linear Gaussians?effective_einstein_radiuson thevis_lpmass?). - Decide whether the global NaN mask should sacrifice the offending latent and keep the rest
instead of dropping the whole block.
Why it matters
Any producer pointed at vis_lp (lens_mass.py --search_name=vis_lp is a supported flag) gets
silently blank latent columns for this reason, unrelated to the retired latent. prefix bug, and
the SED chain seeds sersic_lens_model from vis_lp, so this needs a verdict before the
euclid_dr1_prelim SED fits are trusted.
Related: draft/bug/euclid/catalogue_latent_prefix_blank_columns.md (the retired latent. prefix
this is distinct from) and draft/bug/autofit/aggregate_csv_latent_sigma3_and_silent_none.md (the
silent-None path that hides both).
Activity
Merge pull request #734 from PyAutoLabs/feature/vis-lp-latent-jit-ass…
Jammy2211 commentedon Sep 10, 2026
Shipped
latent_instance_frombuilds the latent instance theFitnessway under JAX (assertions skipped,xpthreaded); NumPy path unchanged. Labelledpending-release.latent_samples_fromnow logs the count and first traceback of the exceptions it swallows, and names the every-sample case. Labelledpending-release.Root cause: the engine's per-sample
jax.jitcould not trace the default assertion check ininstance_from_vector; thevis_lpmodel gained an ordering assertion on its two MGE bases on 2026-09-08, so every sample raisedTracerBoolConversionErrorand was masked to a NaN row. The global NaN mask was not at fault.Remaining, filed in PyAutoMind as
draft/bug/euclid/vis_lp_latent_euclid_jit_trace.md: the euclid pipeline'sLatentEuclid.variablesmust use the new helper, and its uniform source-flux grid must be built on the static NumPy mask (a second, independent trace failure). Then avis_lprerun on tile 102005065 with refreshed RAL libraries is the witness.Generated by Claude Code
fix(latent): build the latent instance the Fitness way under JAX (PyA…