Skip to content

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

Description

@Jammy2211

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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions