From c204082deee5e2481be28b735e42766c8976c0ae Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 10 Sep 2026 18:22:15 +0000 Subject: [PATCH] fix(latent): build the latent instance the Fitness way under JAX (#732) `LatentLens.variables` built its per-sample instance with `model.instance_from_vector(vector=parameters)`, whose default assertion check applies a Python `not` to a traced boolean. The latent engine evaluates `variables` under a per-sample `jax.jit`, so a model carrying any assertion (the ordered two-basis MGE lens light of the Euclid `vis_lp` stage) raised `TracerBoolConversionError` on every sample; the engine swallowed each raise into a NaN row, the global mask found nothing, and the search wrote no latent output at all (RAL job 342398, all ten euclid_dr1_prelim tiles). Add `latent_instance_from(model, parameters, xp)`: the NumPy path is unchanged (assertions checked, a violation raises FitException and is dropped as a NaN row); under JAX it skips assertions and threads `xp`, exactly as `Fitness` does on its JAX path. `LatentLens.variables` uses it and the class docstring tells subclasses that build their own instance to do the same. Tests: the helper on both backends (the default path is shown to raise under jit), and `LatentLens.variables` end to end under `jax.jit` on a model with an assertion, finite and equal to the NumPy evaluation. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01FFSFgEBiZc7kTRutfK1RsN --- autolens/analysis/latent.py | 36 ++++++++- test_autolens/analysis/test_latent.py | 104 ++++++++++++++++++++++++++ 2 files changed, 138 insertions(+), 2 deletions(-) diff --git a/autolens/analysis/latent.py b/autolens/analysis/latent.py index 3b3afcd1f..d0024aafc 100644 --- a/autolens/analysis/latent.py +++ b/autolens/analysis/latent.py @@ -439,6 +439,34 @@ def effective_einstein_radius(fit, magzero, xp=np): return xp.nan +def latent_instance_from(model, parameters, xp=np): + """ + Build the model instance a latent function is evaluated on, on either + backend, the way ``autofit.non_linear.fitness.Fitness`` builds the one the + likelihood is evaluated on. + + On the NumPy path the model's assertions are checked as usual: a sample + that violates one raises ``af.exc.FitException``, which the latent engine + turns into a NaN row and drops. Under JAX that check applies a Python + ``not`` to a traced boolean and raises ``TracerBoolConversionError`` + inside the per-sample ``jax.jit`` the latent batch path uses. The engine + swallowed that raise into a NaN row for **every** sample, so a model with + any assertion (e.g. the ordered two-basis MGE lens light of the Euclid + ``vis_lp`` stage, PyAutoFit#1583) wrote no latent output at all. Under JAX + the assertions are therefore skipped here, as ``Fitness`` skips them + (every sample the search accepted already satisfied them), and ``xp`` is + threaded through so the instance is built on the traced backend. + + Subclasses of :class:`LatentLens` that build their own instance (rather + than composing ``LatentLens.variables``) must go through this helper. + """ + if xp is np: + return model.instance_from_vector(vector=parameters) + return model.instance_from_vector( + vector=parameters, ignore_assertions=True, xp=xp + ) + + LATENT_FUNCTIONS: Dict[str, Callable] = { "total_lens_flux": total_lens_flux, "total_lensed_source_flux": total_lensed_source_flux, @@ -493,7 +521,11 @@ class LatentLens(af.Latent): dispatches the :data:`LATENT_FUNCTIONS` registry on a per-sample fit. Subclass to add project-specific latents (e.g. the Euclid pipeline's aperture fluxes), composing the library values via - ``LatentLens.variables(analysis, parameters, model)``. + ``LatentLens.variables(analysis, parameters, model)``. A subclass that + builds the instance itself must use :func:`latent_instance_from`, which + skips the model's assertions under JAX — checking them inside the + per-sample ``jax.jit`` raises, and the engine turns that raise into a NaN + row for every sample (no latent output at all). """ # The lensing Einstein-radius latent routes through ``ZeroSolver`` @@ -511,7 +543,7 @@ def variables(analysis, parameters, model): if not keys: raise NotImplementedError xp = analysis._xp - instance = model.instance_from_vector(vector=parameters) + instance = latent_instance_from(model=model, parameters=parameters, xp=xp) fit = analysis.fit_from(instance=instance) magzero = analysis.kwargs.get("magzero", None) context = {"fit": fit, "magzero": magzero, "xp": xp} diff --git a/test_autolens/analysis/test_latent.py b/test_autolens/analysis/test_latent.py index 67f2b7e34..3efe41976 100644 --- a/test_autolens/analysis/test_latent.py +++ b/test_autolens/analysis/test_latent.py @@ -870,3 +870,107 @@ def test_analysis_imaging_declares_latent_lens_and_keys_read_config(): "total_source_flux", "total_lens_flux_mujy", ] + + +# --------------------------------------------------------------------------- +# latent_instance_from: assertions on the NumPy path, skipped under JAX +# --------------------------------------------------------------------------- + +def _sersic_model_with_assertion(): + model = af.Model(al.lp.Sersic) + model.add_assertion(model.effective_radius > model.sersic_index) + return model + + +def test_latent_instance_from_numpy_checks_assertions(): + model = _sersic_model_with_assertion() + medians = np.array(model.physical_values_from_prior_medians) + + # sersic_index above effective_radius violates the assertion -> the NumPy + # path raises, exactly as the likelihood path does (Fitness -> resample). + bad = medians.copy() + bad[model.paths.index(("effective_radius",))] = 0.1 + bad[model.paths.index(("sersic_index",))] = 4.0 + with pytest.raises(af.exc.FitException): + _latent_module.latent_instance_from(model=model, parameters=bad, xp=np) + + good = medians.copy() + good[model.paths.index(("effective_radius",))] = 4.0 + good[model.paths.index(("sersic_index",))] = 1.0 + instance = _latent_module.latent_instance_from( + model=model, parameters=good, xp=np + ) + assert instance.effective_radius == pytest.approx(4.0) + + +def test_latent_instance_from_jax_skips_assertions_under_jit(): + """ + The bug: the latent engine evaluates ``Latent.variables`` under a + per-sample ``jax.jit``, and ``instance_from_vector`` with its default + assertion check applies a Python ``not`` to a traced boolean there. The + raise was swallowed into a NaN row for every sample, so a model with any + assertion wrote no latent output at all. + """ + jax = pytest.importorskip("jax") + import jax.numpy as jnp + + model = _sersic_model_with_assertion() + vector = np.array(model.physical_values_from_prior_medians) + vector[model.paths.index(("effective_radius",))] = 4.0 + vector[model.paths.index(("sersic_index",))] = 1.0 + + # The default instance construction cannot be traced. + with pytest.raises(jax.errors.TracerBoolConversionError): + jax.jit(lambda v: model.instance_from_vector(vector=v).effective_radius)( + vector + ) + + # The helper mirrors Fitness's JAX path: assertions skipped, xp threaded. + value = jax.jit( + lambda v: _latent_module.latent_instance_from( + model=model, parameters=v, xp=jnp + ).effective_radius + )(vector) + assert float(value) == pytest.approx(4.0) + + +def test_latent_lens_variables_traces_under_jit_with_model_assertion( + masked_imaging_7x7, +): + """ + End to end: ``LatentLens.variables`` on a model carrying an assertion, + evaluated exactly as the engine's ``jit`` batch path evaluates it, returns + finite values for every enabled latent (instead of raising and being + masked to a NaN row). + """ + jax = pytest.importorskip("jax") + + lens = af.Model(al.Galaxy, redshift=0.5, light=af.Model(al.lp.Sersic)) + source = af.Model(al.Galaxy, redshift=1.0, light=af.Model(al.lp.Sersic)) + model = af.Collection(galaxies=af.Collection(lens=lens, source=source)) + model.add_assertion( + model.galaxies.lens.light.intensity > model.galaxies.source.light.intensity + ) + + vector = np.array(model.physical_values_from_prior_medians) + vector[model.paths.index(("galaxies", "lens", "light", "intensity"))] = 0.1 + vector[model.paths.index(("galaxies", "source", "light", "intensity"))] = 0.05 + + analysis = al.AnalysisImaging(dataset=masked_imaging_7x7, use_jax=True, magzero=25.0) + keys = LatentLens.keys(analysis) + + values = jax.jit( + lambda v: LatentLens.variables(analysis, parameters=v, model=model) + )(vector) + + assert len(values) == len(keys) + assert all(np.isfinite(float(v)) for v in values) + + # And they agree with the NumPy evaluation of the same sample. + analysis_np = al.AnalysisImaging( + dataset=masked_imaging_7x7, use_jax=False, magzero=25.0 + ) + values_np = LatentLens.variables(analysis_np, parameters=vector, model=model) + assert np.allclose( + [float(v) for v in values], [float(v) for v in values_np], rtol=1e-6 + )