Skip to content

chore(config): release profile should not disable JAX — delete the override-list trap#47

Merged
Jammy2211 merged 1 commit into
mainfrom
feature/release-profile-jax-default
Jul 15, 2026
Merged

chore(config): release profile should not disable JAX — delete the override-list trap#47
Jammy2211 merged 1 commit into
mainfrom
feature/release-profile-jax-default

Conversation

@Jammy2211

Copy link
Copy Markdown
Collaborator

Delete the trap instead of documenting it

Five nights of red nightly-release (#44 / PR #45) came down to one thing: a JAX-native script landed without an entry in this profile's hand-maintained PYAUTO_DISABLE_JAX override list, and nothing said so — it silently validated numpy.

PR #45 fixed the symptom by adding the three missing entries. This removes the reason the list exists.

The setting is pure downside in this repo

PYAUTO_DISABLE_JAX is read in exactly one place that matters here: af.Analysis.__init__, where it forces use_jax=False. Every analysis in this repo is af.ex.Analysis or a local subclass, and af.Analysis defaults to use_jax=False.

So the "1" default:

  • is a no-op for the ~35 scripts that never opt in — they run numpy either way, so it buys zero speed;
  • silently sabotages the ~7 that explicitly pass use_jax=True — precisely the scripts written to exercise JAX.

Its only real effect was to require a "remember to add your script" list. That list is the trap: forgetting it does not fail loudly, it passes vacuously. searches/Dynesty_jax and searches/Nautilus_jax did exactly that for as long as they have existed. searches/MultiStartAdam only failed loudly because it happens to carry its own backend guard.

The sibling autofit_workspace release profile has always pinned "0""JAX enabled (release fidelity, not smoke)". This repo was the inconsistent one.

Change

  • PYAUTO_DISABLE_JAX: "1""0" in config/build/env_vars_release.yaml
  • overrides:[] — all five JAX overrides (jax_assertions/, BlackJAXNUTS, MultiStartAdam, Dynesty_jax, Nautilus_jax) become no-ops and are deleted
  • The reasoning is recorded in the file, including an explicit do-not-copy-this note for autolens_workspace_test / autogalaxy_workspace_test, where PyAutoGalaxy/PyAutoLens analyses default to use_jax=True and "1" therefore is not a no-op.

Net: -31 lines of config, and a class of bug that cannot recur, because there is nothing left to forget.

Verification

Resolved every scripts/**/*.py through autobuild.env_config.build_env_for_script under the old and new profiles: 31 scripts see the env value change 10, 11 unchanged. Env value changing is not behaviour changing — for a script that never passes use_jax, af.Analysis defaults False and the var is inert. That claim was then tested, not asserted:

check result
features/latent_nan_robustness.py — opts in, so its JAX half runs for the first time exit 0 — no bug was hiding behind the vacuous pass
searches/Nautilus.pyDISABLE_JAX=1 vs 0 exit 0 / exit 0 — identical
searches/Emcee.pyDISABLE_JAX=1 vs 0 exit 0 / exit 0 — identical
features/latent.pyDISABLE_JAX=1 vs 0 exit 0 / exit 0 — identical
model_composition/model_composition.pyDISABLE_JAX=1 vs 0 exit 0 / exit 0 — identical
MultiStartAdam / Dynesty_jax / Nautilus_jax (already 0 via #45, unchanged here) exit 0 (verified in #44)

Why this rather than the alternatives

Three other levers were investigated today and rejected on evidence:

  1. Promote into smoke_tests.txt — would not have caught it. run_smoke.py hardcodes ENV_VARS_FILE = config/build/env_vars.yaml; the per-PR gate never reads the release profile.
  2. A static config guard ("every script with use_jax=True must resolve DISABLE_JAX=0") — the rule false-positives on simulator_use_jax_parity.py, which uses SimulatorImaging (does not read the env var).
  3. A loud warning in af.Analysis (PyAutoFit#1372) — built, then killed by adversarial review: PyAutoLens flips use_jax before super().__init__(), so the warning never fires for 54 autolens scripts; and PyAutoGalaxy/PyAutoLens analyses default use_jax=True, so it fires on default constructions claiming JAX "was requested" when it was not. Silent where it mattered, noisy where it did not.

The deciding argument: a memory note recording this exact two-profile trap already existed, written the day before, and the trap was still walked into this morning. Documentation did not prevent it. Removing the thing to remember does.

Follow-ups (not here)

  • PyAutoFit#1372 — the real finding under the warning: use_jax/PYAUTO_DISABLE_JAX is handled inconsistently across three layers (af default False+reads env; ag default True+does not; al default True+reads env early, duplicating af). The autolens early flip is a second instance of the same silent-override bug one layer up.
  • autolens_workspace_test / autogalaxy_workspace_test release profiles pin DISABLE_JAX: "1" where analyses default use_jax=True — so their mode=release runs validate numpy for any script not matching a jax_*/ override. Whether that is intended is worth its own look.

PYAUTO_DISABLE_JAX is read in exactly one place that matters to this
workspace: af.Analysis.__init__, where it forces use_jax=False. Every analysis
here is af.ex.Analysis or a local subclass, and af.Analysis defaults to
use_jax=False.

So the "1" default was a no-op for the ~35 scripts that never opt in — they run
numpy either way — and silently sabotaged the ~7 that explicitly pass
use_jax=True, which are exactly the scripts written to exercise JAX. It bought
no speed and cost correctness.

Its only real effect was to require a hand-maintained "remember to add your
script" override list. That list is the trap: a JAX-native script that forgets
it does not fail loudly, it silently validates numpy and passes. Dynesty_jax
and Nautilus_jax did that for as long as they have existed; MultiStartAdam only
failed loudly (nightly-release red five nights, #44) because it happens to
carry its own backend guard.

Pin "0" and delete all five overrides, which become no-ops. This matches the
sibling autofit_workspace release profile, which has always pinned "0" for the
same release-fidelity reason — this repo was the inconsistent one. PR #45 fixed
the symptom by adding the three missing entries; this removes the reason the
list exists.

The file now carries the reasoning, including an explicit do-not-copy note for
autolens_workspace_test / autogalaxy_workspace_test, whose analyses come from
PyAutoGalaxy/PyAutoLens and default to use_jax=True — there "1" is NOT a no-op.

Verified rather than asserted. Resolved every scripts/**/*.py through
autobuild.env_config.build_env_for_script under both profiles: 31 change 1->0,
11 unchanged. A changed env value is not changed behaviour, so:

  features/latent_nan_robustness.py  (opts in; JAX half runs for the first
                                      time)          -> exit 0, no hidden bug
  searches/Nautilus.py               DISABLE_JAX=1/0 -> exit 0 / exit 0
  searches/Emcee.py                  DISABLE_JAX=1/0 -> exit 0 / exit 0
  features/latent.py                 DISABLE_JAX=1/0 -> exit 0 / exit 0
  model_composition/model_composition.py       1/0   -> exit 0 / exit 0

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
@Jammy2211 Jammy2211 added the pending-release PR queued for the next release build label Jul 15, 2026
@Jammy2211
Jammy2211 merged commit e97d71e into main Jul 15, 2026
4 checks passed
@Jammy2211
Jammy2211 deleted the feature/release-profile-jax-default branch July 15, 2026 09:01
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

pending-release PR queued for the next release build

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant