fix: PYAUTO_DISABLE_JAX has exactly one reader (#181 step 7)

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 = (

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

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
 

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
 

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,
 ):

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(