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(