chore(config): release profile should not disable JAX — delete the override-list trap#47
Merged
Merged
Conversation
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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-maintainedPYAUTO_DISABLE_JAXoverride 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_JAXis read in exactly one place that matters here:af.Analysis.__init__, where it forcesuse_jax=False. Every analysis in this repo isaf.ex.Analysisor a local subclass, andaf.Analysisdefaults touse_jax=False.So the
"1"default: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_jaxandsearches/Nautilus_jaxdid exactly that for as long as they have existed.searches/MultiStartAdamonly failed loudly because it happens to carry its own backend guard.The sibling
autofit_workspacerelease profile has always pinned"0"— "JAX enabled (release fidelity, not smoke)". This repo was the inconsistent one.Change
PYAUTO_DISABLE_JAX: "1"→"0"inconfig/build/env_vars_release.yamloverrides:→[]— all five JAX overrides (jax_assertions/,BlackJAXNUTS,MultiStartAdam,Dynesty_jax,Nautilus_jax) become no-ops and are deletedautolens_workspace_test/autogalaxy_workspace_test, where PyAutoGalaxy/PyAutoLens analyses default touse_jax=Trueand"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/**/*.pythroughautobuild.env_config.build_env_for_scriptunder the old and new profiles: 31 scripts see the env value change1→0, 11 unchanged. Env value changing is not behaviour changing — for a script that never passesuse_jax,af.AnalysisdefaultsFalseand the var is inert. That claim was then tested, not asserted:features/latent_nan_robustness.py— opts in, so its JAX half runs for the first timesearches/Nautilus.py—DISABLE_JAX=1vs0searches/Emcee.py—DISABLE_JAX=1vs0features/latent.py—DISABLE_JAX=1vs0model_composition/model_composition.py—DISABLE_JAX=1vs0MultiStartAdam/Dynesty_jax/Nautilus_jax(already0via #45, unchanged here)Why this rather than the alternatives
Three other levers were investigated today and rejected on evidence:
smoke_tests.txt— would not have caught it.run_smoke.pyhardcodesENV_VARS_FILE = config/build/env_vars.yaml; the per-PR gate never reads the release profile.use_jax=Truemust resolveDISABLE_JAX=0") — the rule false-positives onsimulator_use_jax_parity.py, which usesSimulatorImaging(does not read the env var).af.Analysis(PyAutoFit#1372) — built, then killed by adversarial review: PyAutoLens flipsuse_jaxbeforesuper().__init__(), so the warning never fires for 54 autolens scripts; and PyAutoGalaxy/PyAutoLens analyses defaultuse_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)
use_jax/PYAUTO_DISABLE_JAXis handled inconsistently across three layers (af defaultFalse+reads env; ag defaultTrue+does not; al defaultTrue+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_testrelease profiles pinDISABLE_JAX: "1"where analyses defaultuse_jax=True— so their mode=release runs validate numpy for any script not matching ajax_*/override. Whether that is intended is worth its own look.