From 83016c1ea027f785d785aef47f997f99ec9d3606 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 23 Jul 2026 18:35:24 +0100 Subject: [PATCH] fix: PYAUTO_DISABLE_JAX has exactly one reader (#181 step 7) af.Analysis.__init__ is the single resolver of PYAUTO_DISABLE_JAX + jax-availability; AnalysisLens.__init__ now always receives the base-resolved self._use_jax instead of the raw parameter, at all three call sites. Deletes the duplicate env read in AnalysisDataset (design failure mode 8) and FIXES a live bug: AnalysisPoint and the weak analysis silently undid the env downgrade (base set False, AnalysisLens overwrote True), so PYAUTO_DISABLE_JAX=1 was a no-op for them. Construction-level tests: env downgrade respected for AnalysisImaging and AnalysisPoint (the point test fails on pre-fix source), no over-downgrade when the env var is unset. Co-Authored-By: Claude Fable 5 --- autolens/analysis/analysis/dataset.py | 12 +++--- autolens/analysis/analysis/lens.py | 8 ++-- autolens/point/model/analysis.py | 5 ++- autolens/weak/model/analysis.py | 5 ++- .../analysis/test_analysis_dataset.py | 32 ++++++++++++++ .../point/model/test_analysis_point.py | 43 +++++++++++++++++++ 6 files changed, 93 insertions(+), 12 deletions(-) diff --git a/autolens/analysis/analysis/dataset.py b/autolens/analysis/analysis/dataset.py index a71b236e8..4198ab303 100644 --- a/autolens/analysis/analysis/dataset.py +++ b/autolens/analysis/analysis/dataset.py @@ -14,7 +14,6 @@ """ import logging import numpy as np -import os from typing import List, Optional from autonerves import conf @@ -84,11 +83,6 @@ def __init__( anyway. """ - import os - - if os.environ.get("PYAUTO_DISABLE_JAX") == "1": - use_jax = False - super().__init__( dataset=dataset, adapt_images=adapt_images, @@ -99,11 +93,15 @@ def __init__( **kwargs, ) + # `super().__init__` routes through `af.Analysis.__init__`, the single + # reader of the disable-jax env var and the jax-availability check, which + # resolves `self._use_jax`. Forward that resolved value so `AnalysisLens` + # never overwrites it with the raw parameter. AnalysisLens.__init__( self=self, positions_likelihood_list=positions_likelihood_list, cosmology=cosmology, - use_jax=use_jax, + use_jax=self._use_jax, ) self.raise_inversion_positions_likelihood_exception = ( diff --git a/autolens/analysis/analysis/lens.py b/autolens/analysis/analysis/lens.py index 732073dcb..03d558f85 100644 --- a/autolens/analysis/analysis/lens.py +++ b/autolens/analysis/analysis/lens.py @@ -61,9 +61,11 @@ def __init__( self.cosmology = cosmology or Planck15() self.positions_likelihood_list = positions_likelihood_list - # Mirror the autofit Analysis fallback: if jax isn't installed, - # downgrade silently here too (the parent Analysis.__init__ already - # emitted the loud banner — no need to repeat it). + # `use_jax` is expected to already be the base-resolved value + # (`self._use_jax` set by `af.Analysis.__init__`, the single reader of + # the disable-jax env var and the jax-availability check). This guard is + # a defensive, idempotent re-check of jax availability only — it never + # re-reads the env var and does not repeat the parent's loud banner. import importlib.util if use_jax and importlib.util.find_spec("jax") is None: use_jax = False diff --git a/autolens/point/model/analysis.py b/autolens/point/model/analysis.py index 9124a3994..c3da8bb6e 100644 --- a/autolens/point/model/analysis.py +++ b/autolens/point/model/analysis.py @@ -81,7 +81,10 @@ def __init__( """ super().__init__(cosmology=cosmology, use_jax=use_jax, **kwargs) - AnalysisLens.__init__(self=self, cosmology=cosmology, use_jax=use_jax) + # `super().__init__` (af.Analysis) is the single reader of the + # disable-jax env var + the jax-availability check; forward the + # resolved `self._use_jax` so `AnalysisLens` does not overwrite it. + AnalysisLens.__init__(self=self, cosmology=cosmology, use_jax=self._use_jax) self.dataset = dataset diff --git a/autolens/weak/model/analysis.py b/autolens/weak/model/analysis.py index bcc6f9f35..dff3a61c0 100644 --- a/autolens/weak/model/analysis.py +++ b/autolens/weak/model/analysis.py @@ -68,7 +68,10 @@ def __init__( """ super().__init__(cosmology=cosmology, use_jax=use_jax, **kwargs) - AnalysisLens.__init__(self=self, cosmology=cosmology, use_jax=use_jax) + # `super().__init__` (af.Analysis) is the single reader of the + # disable-jax env var + the jax-availability check; forward the + # resolved `self._use_jax` so `AnalysisLens` does not overwrite it. + AnalysisLens.__init__(self=self, cosmology=cosmology, use_jax=self._use_jax) self.dataset = dataset diff --git a/test_autolens/analysis/analysis/test_analysis_dataset.py b/test_autolens/analysis/analysis/test_analysis_dataset.py index 4767169f2..04d125e5d 100644 --- a/test_autolens/analysis/analysis/test_analysis_dataset.py +++ b/test_autolens/analysis/analysis/test_analysis_dataset.py @@ -1,4 +1,5 @@ from pathlib import Path +import importlib.util import os import pytest @@ -12,6 +13,37 @@ directory = Path(__file__).resolve().parent +def _jax_installed() -> bool: + return importlib.util.find_spec("jax") is not None + + +def test__pyauto_disable_jax_env_downgrades_use_jax__imaging( + monkeypatch, masked_imaging_7x7 +): + # Regression cover for the deleted local env read in `AnalysisDataset`: + # the disable-jax env var must still downgrade `use_jax`, now resolved + # solely by `af.Analysis.__init__` (the single reader) and forwarded to + # `AnalysisLens` as `self._use_jax`. + monkeypatch.setenv("PYAUTO_DISABLE_JAX", "1") + + analysis = al.AnalysisImaging(dataset=masked_imaging_7x7, use_jax=True) + + assert analysis._use_jax is False + + +@pytest.mark.skipif(not _jax_installed(), reason="jax not installed") +def test__use_jax_true_env_unset__not_downgraded__imaging( + monkeypatch, masked_imaging_7x7 +): + # No over-downgrade: with the env var unset and jax installed, + # `use_jax=True` must survive as `self._use_jax is True`. + monkeypatch.delenv("PYAUTO_DISABLE_JAX", raising=False) + + analysis = al.AnalysisImaging(dataset=masked_imaging_7x7, use_jax=True) + + assert analysis._use_jax is True + + def test__modify_before_fit__inversion_no_positions_likelihood__raises_exception( masked_imaging_7x7, ): diff --git a/test_autolens/point/model/test_analysis_point.py b/test_autolens/point/model/test_analysis_point.py index 207cb880a..255246f9a 100644 --- a/test_autolens/point/model/test_analysis_point.py +++ b/test_autolens/point/model/test_analysis_point.py @@ -1,4 +1,7 @@ from pathlib import Path +import importlib.util + +import pytest import autofit as af import autolens as al @@ -8,6 +11,46 @@ directory = Path(__file__).resolve().parent +def _jax_installed() -> bool: + return importlib.util.find_spec("jax") is not None + + +def test__pyauto_disable_jax_env_downgrades_use_jax__point( + monkeypatch, point_dataset +): + # THE BUG TEST. Before the one-reader fix `AnalysisPoint` had no local + # env read and `AnalysisLens.__init__` overwrote the base-resolved + # `self._use_jax` with the raw `use_jax` parameter, so the disable-jax + # env var was silently a no-op (base set False, AnalysisLens set True). + # It must now downgrade to False. + monkeypatch.setenv("PYAUTO_DISABLE_JAX", "1") + + solver = al.m.MockPointSolver(model_positions=point_dataset.positions) + + analysis = al.AnalysisPoint( + dataset=point_dataset, solver=solver, use_jax=True + ) + + assert analysis._use_jax is False + + +@pytest.mark.skipif(not _jax_installed(), reason="jax not installed") +def test__use_jax_true_env_unset__not_downgraded__point( + monkeypatch, point_dataset +): + # No over-downgrade: with the env var unset and jax installed, + # `use_jax=True` must survive as `self._use_jax is True`. + monkeypatch.delenv("PYAUTO_DISABLE_JAX", raising=False) + + solver = al.m.MockPointSolver(model_positions=point_dataset.positions) + + analysis = al.AnalysisPoint( + dataset=point_dataset, solver=solver, use_jax=True + ) + + assert analysis._use_jax is True + + def _test__make_result__result_imaging_is_returned(point_dataset): model = af.Collection( galaxies=af.Collection(