fix: searches shard red — missing __Env__ declarations + optax/blackjax never installed

[Jammy2211]

Overview

PyAutoHeart/workspace-validation has been red since 2026-07-26 on the
run_scripts (3.12, autofit_test, searches) shard — four failures with two
independent root causes, both confirmed by local reproduction. One is a
missing test-harness declaration; the other is a real user-facing packaging
gap (pip install autolens[jax] does not install optax, so
af.MultiStartAdam / af.MultiStartProdigy raise ImportError for users).

scripts/searches/BlackJAXNUTS.py        FAIL  ModuleNotFoundError: No module named 'blackjax'
scripts/searches/MultiStartAdam.py      FAIL  ImportError: requires optional `jax` and `optax`
scripts/searches/MultiStartProdigy.py   FAIL  AssertionError: 1.0
scripts/searches/MultiStartResurrect.py FAIL  KeyError: 'n_resurrections'

Plan

• (a) MultiStartProdigy.py and MultiStartResurrect.py carry no __Env__
  declaration, so they inherit the smoke defaults PYAUTO_TEST_MODE=2 +
  PYAUTO_DISABLE_JAX=1 and run the bypass path instead of a real JAX search.
  Add the missing declaration to both.
• (b) The validation smoke leg installs autolens[optional], whose extras
  chain stops at autonerves[jax] and never reaches autofit[jax] (optax) or
  autofit[optional] (blackjax). Fix the chain in PyAutoGalaxy so optax follows
  the jax extra everywhere, and install autofit[optional] in the workflow for
  blackjax.
• Audit the sibling JAX search scripts and report (do not silently fix).

Detailed implementation plan

Affected Repositories

• autofit_workspace_test (primary)
• PyAutoGalaxy
• PyAutoHeart

Branch Survey

Repository	Current Branch	Dirty?
./autofit_workspace_test	main	clean
./PyAutoGalaxy	main	clean
./PyAutoHeart	main	clean

Suggested branch: feature/validation-searches-env-optax

Reproduction (local, byte-identical to CI)

cd autofit_workspace_test
PYAUTO_TEST_MODE=2 PYAUTO_DISABLE_JAX=1 PYAUTO_SKIP_WORKSPACE_VERSION_CHECK=1 \
  python3 scripts/searches/MultiStartProdigy.py
#   Recovered: centre=50.000, normalization=1.000, sigma=15.000
#   AssertionError: 1.0

1.0 is the LogUniform(1e-2, 1e2) midpoint the bypass returns for
normalization (truth 25.0). n_resurrections is absent because the bypass
never builds a search_internal. Both scripts pass on a normal run.

Cause: the #187/#189 ENV migration (2026-07-23) named only the two siblings
(MultiStartAdam, BlackJAXNUTS) and missed these two.

Implementation Steps

1. autofit_workspace_test — append an __Env__ section carrying
   ENV: real_search jax to the end of the module docstring of
   scripts/searches/MultiStartProdigy.py and
   scripts/searches/MultiStartResurrect.py. Mirror
   scripts/searches/MultiStartAdam.py:24-30 exactly: column-0 __Env__
   header, exactly one ENV: line, no leading # — the legacy comment form
   now raises (PyAutoHands/autohands/env_config.py:82-92). The tokens release
   PYAUTO_TEST_MODE and PYAUTO_DISABLE_JAX (env_config.py:43-54).

2. PyAutoGalaxy — pyproject.toml jax extra: "autonerves[jax]" →
   "autofit[jax]". Strictly additive (autofit[jax] = ["autonerves[jax]", "optax>=0.2.5"]), and PyAutoGalaxy already depends on autofit
   (pyproject.toml:29), so no new dependency edge. autolens[jax] → autogalaxy[jax] inherits it.

3. PyAutoHeart — .github/workflows/workspace-validation.yml, smoke
   install step (~236-248): add pip install "autofit[optional]" for blackjax,
   commented in the style of the existing nufftax precedent directly above.
   mode=release already installs autofit[optional]==$TESTPYPI_VERSION, so
   only the smoke leg changes.

4. Audit / report only — Nautilus_jax.py and Dynesty_jax.py also use JAX
   with no declaration. They do not fail (no truth-recovery assert) but are
   therefore running the numpy path under smoke, contributing zero JAX
   coverage. Report and recommend; do not fold in silently. The remaining
   searches scripts are numpy-only and correctly need nothing.

Key Files

• scripts/searches/MultiStartProdigy.py — add __Env__ section
• scripts/searches/MultiStartResurrect.py — add __Env__ section
• scripts/searches/MultiStartAdam.py — the reference form to mirror
• PyAutoGalaxy/pyproject.toml — jax extra
• PyAutoHeart/.github/workflows/workspace-validation.yml — smoke install step

Verification

• Both scripts pass under the exact smoke defaults that currently fail.
• PyAutoHands/autohands/validate_env_profiles.py reports both as declared, and
  the resolved-env diff shows the two vars released for them and unchanged
  for every other script (config changes are verified by resolved-env diff, not
  smoke alone).
• Scratch venv: pip install -e PyAutoGalaxy[jax] then import optax.
• workspace-validation green on the next scheduled run — read the verdict via
  the Actions API (gh pr checks is unparseable on this gh build). Do
  not hand-dispatch the nightly release.

Notes

• Brain phase-split overridden. pyauto-brain feature scored this
  too-large (score 20) and proposed a 4-phase split. That score tracks the repo
  count (4), not change size — the total diff is two docstring sections, one
  pyproject.toml line and one pip install line, with no API change and no
  public-API ripple. Kept as one easy task; one PR per repo.
• The morning /wake_up digest first attributed these failures to the MultiStart
  cadence PRs (#1421/#1423). That was wrong — they merged 2026-07-27 16:00 BST,
  ~10h after the 06:15 failing run.
• Do not "fix" BlackJAXNUTS.py by deleting it or dropping blackjax —
  blackjax is a kept dependency.

Original Prompt

Click to expand starting prompt

Original request:

can we sequentially tackle everything in 3 overnight jobs red

(from the 2026-07-27 /wake_up digest; this is the first of the three red jobs.)

[Jammy2211]

Library PRs created

All three fixes are open and labelled pending-release:

Repo	PR	Fixes
autofit_workspace_test	#78	missing __Env__ declarations (cause a)
PyAutoGalaxy	#530	optax never reaching the jax chain (cause b)
PyAutoHeart	#111	blackjax absent from the smoke leg (cause b)

All three are needed for a green autofit_test/searches shard.

Verification done

• Resolved-env diff, all 63 scripts: exactly 2 changed, each releasing only PYAUTO_TEST_MODE + PYAUTO_DISABLE_JAX; the other 61 byte-identical. This is the authoritative check for a config change — smoke alone would not prove the absence of collateral effects.
• Both scripts now run the real JAX search and pass:
  • MultiStartProdigy → centre=50.154, normalization=25.195, sigma=9.858
  • MultiStartResurrect → n_resurrections 0 (off) / 1 (on), identical basin, np.allclose safety property holds
• PyAutoGalaxy built-wheel metadata: Requires-Dist: autofit[jax]; extra == "jax" — chain to optax>=0.2.5 closed.
• workspace-validation.yml parses; only the mode=smoke step changed.

Heart gate

YELLOW (score 65), no RED — acknowledged by the human. None of the yellow reasons is caused by this change, and the first (workspace validation not passing (13 failed)) is what this fixes.

Audit finding — reported, not fixed

Nautilus_jax.py and Dynesty_jax.py also carry no ENV declaration. They do not fail, because neither asserts truth recovery — but analysis.py:68 silently forces use_jax=False when PYAUTO_DISABLE_JAX=1, and use_jax=True is the only thing separating them from their Nautilus.py / DynestyStatic.py siblings. So under smoke they are duplicate numpy runs contributing zero JAX coverage.

Left out of this PR deliberately — it changes what those scripts exercise (and their runtime), which is a separate call. Recommend adding ENV: jax to both as a follow-up.

Correction

The morning /wake_up digest attributed these failures to the MultiStart cadence PRs (#1421/#1423). That was wrong — they merged 2026-07-27 16:00 BST, ~10h after the 06:15 failing run. The real trigger is the #187/#189 ENV migration (2026-07-23).

[Jammy2211]

Shipped — all three PRs merged 2026-07-27: autofit_workspace_test#78, PyAutoLabs/PyAutoGalaxy#530, PyAutoLabs/PyAutoHeart#111. __Env__ declarations added to MultiStartProdigy/MultiStartResurrect (real_search+jax), optax now follows the jax extra, blackjax installed via autofit[optional] in the smoke leg. Completion recorded in PyAutoMind (merged by the #657 session to clear the PyAutoGalaxy worktree claim; local worktree cleanup deferred to the owning session).