From 43088048a8121aa9e034ec1e98f591bb8ed650bb Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 15 Jul 2026 09:54:23 +0100 Subject: [PATCH] chore(config): stop the release profile disabling JAX (#46) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- config/build/env_vars_release.yaml | 62 +++++++++++++++--------------- 1 file changed, 30 insertions(+), 32 deletions(-) diff --git a/config/build/env_vars_release.yaml b/config/build/env_vars_release.yaml index aa630f6..5480e37 100644 --- a/config/build/env_vars_release.yaml +++ b/config/build/env_vars_release.yaml @@ -29,40 +29,38 @@ defaults: PYAUTO_TEST_MODE: "0" # real searches — no reduced-iteration bypass - PYAUTO_DISABLE_JAX: "1" # force use_jax=False by default; jax_assertions/ + searches/BlackJAXNUTS re-enable below + PYAUTO_DISABLE_JAX: "0" # JAX enabled (release fidelity, not smoke) — see "Why 0" below PYAUTO_SKIP_WORKSPACE_VERSION_CHECK: "1" # TestPyPI dev version won't match the workspace pin JAX_ENABLE_X64: "True" # enable 64-bit precision in JAX NUMBA_CACHE_DIR: "/tmp/numba_cache" # writable cache dir for numba MPLCONFIGDIR: "/tmp/matplotlib" # writable config dir for matplotlib -overrides: - # Aggregator / database / minimal-output / features/assertion scripts all - # need a real sampler run — PYAUTO_TEST_MODE="0" above already satisfies - # that for every one of them, so unlike env_vars.yaml (whose smoke default - # bypasses the sampler entirely), no override is needed here. - # - # JAX assertion scripts exist specifically to test JAX behavior — disabling - # JAX makes their assertions vacuous (Analysis._use_jax_for_visualization is - # silently flipped to False, so cached attrs like _jitted_fit_from never get - # set, and downstream assertions about those attrs hit __getattr__). - - pattern: "jax_assertions/" - set: { PYAUTO_DISABLE_JAX: "0" } - # BlackJAXNUTS is a gradient-based MCMC that requires JAX (jax.grad of the - # autofit Fitness) — needs JAX enabled. PYAUTO_TEST_MODE="0" (real sampler) - # already satisfied by defaults. - - pattern: "searches/BlackJAXNUTS" - set: { PYAUTO_DISABLE_JAX: "0" } - # MultiStartAdam is a JAX-native gradient MAP search (jax.vmap of - # jax.value_and_grad over the autofit Fitness). Unlike the samplers below it - # hard-raises when the analysis is not on the JAX backend, so the disabling - # default turned it into a release-validation failure the moment it landed. - - pattern: "searches/MultiStartAdam" - set: { PYAUTO_DISABLE_JAX: "0" } - # The *_jax sampler scripts exist to validate the JAX likelihood path. They - # carry no hard backend guard, so under the disabling default they silently - # validated the NumPy path and passed vacuously — the same failure mode the - # jax_assertions/ note above describes, minus the loud failure. - - pattern: "searches/Dynesty_jax" - set: { PYAUTO_DISABLE_JAX: "0" } - - pattern: "searches/Nautilus_jax" - set: { PYAUTO_DISABLE_JAX: "0" } +# Why PYAUTO_DISABLE_JAX is "0" here, and why there is no per-script JAX +# override list any more (2026-07-15) +# --------------------------------------------------------------------------- +# 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 in this repo is `af.ex.Analysis` (or a local subclass of it), and +# `af.Analysis` defaults to `use_jax=False`. +# +# So the old "1" default was a no-op for the ~35 scripts that never opt in — +# they run on numpy either way — and silently sabotaged the ~7 that explicitly +# pass `use_jax=True`, which are precisely 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, which is a +# trap: a JAX-native script that forgets it does not fail loudly, it silently +# validates numpy and passes. `searches/Dynesty_jax` and `searches/Nautilus_jax` +# did exactly that for as long as they existed, and `searches/MultiStartAdam` +# only failed loudly (taking nightly-release red for five nights, +# autofit_workspace_test#44) because it happens to carry its own backend guard. +# +# "0" deletes the trap rather than documenting it: there is no list to forget, +# and a script that asks for JAX gets JAX. This also matches the sibling +# `autofit_workspace` release profile, which has always pinned "0" for the same +# release-fidelity reason. +# +# NOTE this reasoning is specific to this repo. Do NOT copy it to +# autolens_workspace_test / autogalaxy_workspace_test: their analyses come from +# PyAutoGalaxy/PyAutoLens, which default to `use_jax=True`, so there "1" really +# does force numpy across the board and is not a no-op. +overrides: []