refactor(analysis): use_jax/PYAUTO_DISABLE_JAX handled inconsistently across af/ag/al

[Jammy2211]

The bug

autofit/non_linear/analysis/analysis.py:67-69:

import os
if os.environ.get("PYAUTO_DISABLE_JAX") == "1":
    use_jax = False

A caller that explicitly passed use_jax=True is silently downgraded to the NumPy path — no warning, no record. Three lines below, the other way use_jax=True can fail to hold (JAX not installed) warns loudly in a formatted box. The env-var override is the only silent one.

Why it matters — it manufactures vacuous tests

PYAUTO_DISABLE_JAX is consumed in exactly one place (this constructor — confirmed by grep across PyAutoConf/PyAutoFit/PyAutoArray). So every workspace script that builds an Analysis(..., use_jax=True) to assert JAX behaviour instead validates the NumPy path, while presenting as a JAX test, and passes.

Observed, not theorised: autofit_workspace_test's searches/Dynesty_jax.py and Nautilus_jax.py were doing exactly this under the mode=release profile (PYAUTO_DISABLE_JAX: "1") and passing vacuously for as long as they have existed. They were only caught because a sibling, MultiStartAdam, happens to carry its own hard backend guard (multi_start_gradient/search.py:120-126) and so failed loudly instead — taking nightly-release red for five nights (autofit_workspace_test#44 / PR #45).

That is the pattern: scripts with a guard fail loudly; scripts without pass vacuously and nobody learns. This silent flip is the root cause of the class.

Plan (high level)

• When the caller explicitly passed use_jax=True and the env var overrides it, say so loudly instead of silently downgrading.
• Mirror the existing JAX-not-installed warning box directly below, so the two failure modes look alike.
• Warn once per process, not once per Analysis.
• No workspace changes here — the scripts this surfaces are a follow-up.

Plan (detailed)

In Analysis.__init__, inside the existing env-var branch: if use_jax is truthy (the parameter defaults to False, so truthiness already means "explicitly requested"), emit a loud warnings.warn mirroring the box below, then downgrade as today. Guard with a module-level flag so one process warns once.

Warn, do not raise. PYAUTO_DISABLE_JAX=1 is legitimate and deliberate — the smoke env profile uses it to run every script fast on NumPy. Raising would break every smoke run of every use_jax=True script. The warning is accurate there rather than noise: it says "this script's JAX assertions are vacuous right now", which is exactly true.

Testing

Unit tests in test_autofit/, NumPy-only (library unit tests never import JAX — and need not here: the warning fires inside the env-var branch, before the JAX-install check):

• use_jax=True + PYAUTO_DISABLE_JAX=1 → warns; _use_jax is False
• use_jax=True + env unset → no warning
• use_jax=False + PYAUTO_DISABLE_JAX=1 → no warning (nothing was overridden)
• two analyses in one process → exactly one warning

Why not the alternatives (considered, rejected 2026-07-15)

• Promote MultiStartAdam into smoke_tests.txt (the question that started this): it would not have caught the bug. run_smoke.py hardcodes ENV_VARS_FILE = config/build/env_vars.yaml, so the per-PR gate never reads the release profile, which is where the missing override lived. It also cuts against "smoke tests are a small curated subset".
• A static per-repo config guard ("every script declaring use_jax=True must resolve DISABLE_JAX=0 in env_vars_release.yaml"): the naive string-matching rule false-positives — autolens_workspace_test/scripts/imaging/simulator_use_jax_parity.py contains use_jax=True but constructs al.SimulatorImaging, which does not read the env var, so its parity assertion is perfectly valid. The correct rule ("constructs an Analysis with use_jax=True") is brittle to detect statically across subclasses. One library change covers all three workspaces and every future script, with no per-repo config.

Blast radius (measured, not estimated)

Resolved autobuild.env_config.build_env_for_script over each release profile, for every script containing use_jax=True:

repo	declares use_jax=True	JAX actually ON	resolves to JAX OFF
autofit_workspace_test	7	6	1
autolens_workspace_test	54	50	4
autogalaxy_workspace_test	36	34	2

The 7 JAX-OFF scripts are what this warning will surface under mode=release. ~2 are false positives (the simulator_use_jax_parity.py pair — SimulatorImaging, unaffected). The rest are genuine and want triage as a follow-up:

• autofit_workspace_test/scripts/features/latent_nan_robustness.py
• autolens_workspace_test/scripts/imaging/visualization.py
• autolens_workspace_test/scripts/latent/latent_nan_robustness.py
• autogalaxy_workspace_test/scripts/imaging/visualization.py
• autogalaxy_workspace_test/scripts/latent/latent_nan_robustness.py

latent_nan_robustness.py deliberately builds both a NumPy and a JAX analysis, so only its JAX half is vacuous — triage it, do not blindly override it.

---

Prompt: PyAutoMind/active/analysis_silently_swallows_use_jax_override.md
Branch: feature/analysis-loud-use-jax-override
Autonomy: supervised (bug cap) — plan recorded here per the autonomy contract; ship sign-off and merge stay human.

[Jammy2211]

⛔ Adversarial review killed this design — NOT shipping it

I asked myself to disprove my own conclusions before shipping. Two of them are false, both proven empirically, and together they invert the fix: it is silent exactly where it matters and loud exactly where it should not be.

✗ False claim 1: "PYAUTO_DISABLE_JAX is consumed in exactly one place"

My original grep covered only PyAutoConf/PyAutoFit/PyAutoArray. Re-run across all libraries, there are two consumers:

• PyAutoFit/autofit/non_linear/analysis/analysis.py:68
• PyAutoLens/autolens/analysis/analysis/dataset.py:89 ← missed

And PyAutoLens flips before delegating:

if os.environ.get("PYAUTO_DISABLE_JAX") == "1":
    use_jax = False          # line 73-74
super().__init__(..., use_jax=use_jax, ...)   # line 76-82  -> af.Analysis sees False

So af.Analysis.__init__ receives use_jax=False, if use_jax: is falsy, and the warning never fires for any autolens analysis — 54 of the 97 scripts, the largest surface. Empirically:

autolens AnalysisImaging(dataset=ds, use_jax=True)  <- EXPLICIT request
  _use_jax          : False (silently flipped to numpy)
  my warning fired? : NO  <-- HOLE: al flips before super()

My end-to-end test passed only because Nautilus_jax.py uses af.ex.Analysis, which reaches the base class directly. I validated on the easy path.

✗ False claim 2: "use_jax defaults to False, so truthiness means the caller explicitly asked"

True for af.Analysis — and false for essentially the entire science surface:

class	use_jax default
af.Analysis	False
ag.AnalysisDataset / AnalysisImaging / AnalysisInterferometer / ellipse	True
al.AnalysisDataset / al.point	True

PyAutoGalaxy never reads the env var, so it passes its default True straight to af.Analysis, which warns. Empirically, on a construction where use_jax was never passed:

autogalaxy AnalysisImaging(dataset=ds)  <- use_jax NEVER passed
  _use_jax          : False
  my warning fired? : YES -- fires on a DEFAULT construction

The message says "use_jax=True was requested". Nobody requested it. Under the smoke profile (DISABLE_JAX=1) this would fire on essentially every autogalaxy/autolens script in every sweep, with a false claim — the textbook way to train everyone to ignore a warning.

✓ Claims that survived

• PYAUTO_DISABLE_JAX really is the only mechanism, and SimulatorImaging really does not read it → simulator_use_jax_parity.py is a genuine false positive of the static-guard rule. Confirmed.
• No filterwarnings = error anywhere (all four libraries use ignore rules only), and no library CI sets PYAUTO_DISABLE_JAX — so the warning would not have broken any downstream suite. That risk was unfounded.
• Dynesty_jax/Nautilus_jax really were vacuous, and the fix really does work for af-level analyses.
• Full pytest test_autofit/ green (1483 passed); the new tests genuinely fail when the fix is disabled.

The real finding underneath

use_jax / PYAUTO_DISABLE_JAX is handled inconsistently across three layers:

layer	default	reads env?
af.Analysis	False	yes
ag.*	True	no
al.AnalysisDataset	True	yes — duplicating af, and early

The autolens early flip is itself a second instance of the silent-override bug, one layer up. "One clean library change covers everything" was wrong.

Where this leaves the options

1. Remove autolens's duplicate early flip — unambiguously correct regardless of what else happens. It is redundant with the base class and it is what hides the override. Care needed: it also feeds AnalysisLens.__init__(use_jax=use_jax), so it must pass self._use_jax after super().__init__() rather than just deleting the branch.
2. Sentinel (use_jax: Optional[bool] = None threaded through af/ag/al) — the only way to make "explicit request" knowable, and therefore the only way the warning can be correct. A real 3-repo API change.
3. Drop the warning; just triage the ~5 genuinely vacuous scripts. Cheapest thing that addresses the actual observed harm.

Parking for a human call rather than guessing. Nothing pushed; the branch is local and uncommitted.

[Jammy2211]

Superseded as the lever — but the finding underneath stands

The warning built here is abandoned (see the adversarial review above: silent for 54 autolens scripts because PyAutoLens flips use_jax before super().__init__(), and noisy-with-a-false-claim on PyAutoGalaxy default constructions because ag.*/al.* default use_jax=True). Nothing from that branch was pushed.

The vacuous-JAX problem that motivated it is instead being deleted at the source, in the workspace: autofit_workspace_test#46 / PR #47 flips that repo's release profile to PYAUTO_DISABLE_JAX: "0" and drops the entire per-script override list. The setting was a no-op for ~35 of 42 scripts there (af.Analysis defaults use_jax=False) and sabotage for the 7 that opt in — so removing it deletes the bug class rather than announcing it.

This issue stays open for the real finding, which is independent of that and still true:

use_jax / PYAUTO_DISABLE_JAX is handled inconsistently across three layers:

layer	use_jax default	reads PYAUTO_DISABLE_JAX?
af.Analysis	False	yes
ag.* (AnalysisDataset, AnalysisImaging, AnalysisInterferometer, ellipse)	True	no
al.AnalysisDataset / al.point	True	yes — duplicating af, and before super().__init__()

Two concrete consequences, both verified:

1. The autolens early flip is itself a silent-override bug, one layer up. It duplicates logic the base class already owns, and it is what makes an explicit use_jax=True vanish without trace. Fixing it is unambiguously right regardless of anything else — but note it also feeds AnalysisLens.__init__(use_jax=use_jax), so the branch cannot simply be deleted: it must pass self._use_jax after super().__init__().
2. "Explicit request" is currently unknowable to the base class, because ag/al always pass a materialised True. Any future logic that wants to distinguish "user asked for JAX" from "defaulted to JAX" needs a sentinel (use_jax: Optional[bool] = None) threaded through all three repos — a real API change worth its own decision.

Retitling this to the inconsistency, which is the durable issue.