fix(latent): build the latent instance the Fitness way under JAX (#732)

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}

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
+    )