feat: promote JAX from optional extra to default dependency

[Jammy2211]

Overview

Promote JAX from the optional [jax] extra to a default dependency across the library stack (autonerves → autofit / autoarray → autogalaxy → autolens). CPU jaxlib wheels exist for every platform we support at requires-python >=3.12 (win_amd64, manylinux x86_64/aarch64, macOS arm64; jaxnnls is pure Python), performance without JAX is unacceptable as a default experience, and promotion eliminates the pip extras-chain history-walk trap (#687) for new users. The NumPy-only code path remains fully supported — and gains, for the first time, an explicit CI leg that tests the jax-not-installed environment.

Plan

• Promote each library's [jax] extra into base dependencies, bottom-up through the chain, bumping intra-family floors to the first promoted version at each layer.
• Declare jax/jaxlib with environment markers so Intel Macs (no jaxlib macosx_x86_64 wheels ≥0.7) resolve cleanly to NumPy-only instead of failing at install. (Human-decided 2026-08-19.)
• Keep [jax] as a no-op alias extra everywhere; widen the jax cap <0.11.0 → <0.12.0.
• Add a no-JAX-installed test job to PyAutoHeart's central reusable lib-tests.yml.
• Add a one-time loud import warning in autonerves when JAX is absent (NumPy-only mode, reduced performance).
• Rewrite install docs: JAX-default is the rule; Intel Mac NumPy-only fallback is the documented exception.

Detailed implementation plan

Affected Repositories

• PyAutoLens (primary)
• PyAutoNerves
• PyAutoFit
• PyAutoArray
• PyAutoGalaxy
• PyAutoHeart (CI workflow only)

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoNerves	main	clean
./PyAutoFit	main	clean
./PyAutoArray	main	clean
./PyAutoGalaxy	main	clean
./PyAutoLens	main	clean
./PyAutoHeart	main	clean

Suggested branch: feature/jax-default-dependency

Note: PyAutoFit is also claimed by stored-sample-reconstruction-guard (#1486, live). Human accepted proceeding anyway — this task touches only PyAutoFit/pyproject.toml, near-zero overlap.

Implementation Steps

1. PyAutoNerves pyproject.toml — move into dependencies:
   • jax>=0.7.0,<0.12.0 ; sys_platform != "darwin" or platform_machine == "arm64"
   • jaxlib>=0.7.0,<0.12.0 ; sys_platform != "darwin" or platform_machine == "arm64"
   • jaxnnls==1.0.1 ; sys_platform != "darwin" or platform_machine == "arm64" (pure Python, but useless without jax)
   • Keep jax = []-style alias extra (or same content) so autonerves[jax] stays valid; update the extras-chain comment block.
2. PyAutoNerves jax_wrapper — one-time warnings.warn (or logger warning) when jax import fails: NumPy-only mode, performance significantly reduced, pointer to install docs. No behavioural change otherwise.
3. PyAutoFit pyproject.toml — add optax>=0.2.5 to dependencies; bump autonerves floor to the first promoted autonerves release; jax = [] alias (keep optax there or empty — alias only).
4. PyAutoArray pyproject.toml — bump autonerves floor; jax extra becomes alias. nufftax / pynufft / tfp-nightly stay in optional (interferometer extras are NOT promoted).
5. PyAutoGalaxy pyproject.toml — add jax_zero_contour>=2.0.0,<3.0.0 to dependencies; bump autofit/autoarray floors; alias extra.
6. PyAutoLens pyproject.toml — bump autogalaxy floor; alias extra.
7. All five repos — update optional extras that currently self-reference [jax] (now redundant but harmless via alias); refresh the "Floor, not a pin" comment blocks to describe the post-promotion rationale.
8. PyAutoHeart .github/workflows/lib-tests.yml — add a unittest-nojax job (Python 3.13 only): identical install, then pip uninstall -y jax jaxlib jaxnnls optax jax_zero_contour, then run the suite. Confirms module import + NumPy paths in a genuinely jax-absent environment (today's env always has jax via [optional]).
9. Docs — PyAutoLens/docs/installation/{pip,conda}.md (+ equivalent pages in the other repos' docs): pip install autolens is now JAX-enabled; remove "[jax] extra recommended" framing; GPU install guidance unchanged; note Intel Macs fall back to NumPy-only automatically (markers) and will see the import warning.
10. Tests — unit suites stay NumPy-only by policy; no fallback-path code changes. Verify each repo's suite passes locally; the new CI leg is the durable guard.

Key Files

• PyAutoNerves/pyproject.toml — jax promotion + markers + cap widen
• PyAutoNerves/autonerves/**/jax_wrapper*.py — no-jax import warning
• PyAutoFit/pyproject.toml, PyAutoArray/pyproject.toml, PyAutoGalaxy/pyproject.toml, PyAutoLens/pyproject.toml — promotion + floors + alias extras
• PyAutoHeart/.github/workflows/lib-tests.yml — no-jax CI leg
• PyAutoLens/docs/installation/pip.md, conda.md — install docs rewrite

Release ordering

Promotion is only effective once released; the floor bumps must walk the chain bottom-up (autonerves first), mirroring the 2026.7.29.2 extras-introduction release. Until then, source installs in this workspace behave identically (jax already installed).

Watch-outs

• Import-time cost: the jax import chain is ~43% of import autolens time; making it default makes every import pay it. Accepted; optimizing import time is out of scope here.
• GPU installs unchanged: CPU jax by default; CUDA remains a manual step.

Original Prompt

Click to expand starting prompt

PyAutoMind prompt: draft/maintenance/libraries/jax_default_dependency.md

Original request (verbatim):

is it time to make JAX not an optional but a default dependency of
autolens? Its easy to install, and performance is crap without it, thoughts?

Ok good, I think autofit and autogalaxy should also be JAX default at this
point, I agree that we should retain the full support for numpy only path,
perhaps we need some sort of explicit test to make sure this path works
cause most environment tests probably install jax? And yes definitely
document how mac users can still install via a no jax route, but this is
now the exception with special pip install options rather than the rule.

See the prompt file for the full scoped body (promotion layout, no-JAX CI leg, docs, watch-outs).

[Jammy2211]

Library PRs Created

All six PRs are open, every test suite green, pending-release label verified on the five library PRs:

Repo	PR	Tests
PyAutoNerves	PyAutoLabs/PyAutoNerves#150	157 passed
PyAutoFit	PyAutoLabs/PyAutoFit#1503	2005 passed, 16 skipped
PyAutoArray	PyAutoLabs/PyAutoArray#450	1062 passed
PyAutoGalaxy	PyAutoLabs/PyAutoGalaxy#574	1111 passed
PyAutoLens	#703	538 passed
PyAutoHeart	PyAutoLabs/PyAutoHeart#150 (no-JAX CI leg; no release label)	484 passed

Extra verification beyond CI: clean-venv install of the promoted autonerves pulled jax 0.11.1 + jaxnnls automatically (markers + widened cap confirmed working); with jax then stripped, the suite passed 157/157 and the new import warning fired exactly once.

API changes: none — packaging, CI, messages, and install docs only. The [jax] extras are kept as no-op aliases, so nothing downstream breaks.

Workspace impact (option ii-lite, docs-only, non-blocking): ~30 files across autolens/autogalaxy/autofit workspaces + HowToLens/HowToGalaxy (plus generated notebook/markdown mirrors) carry now-stale prose recommending the [jax] extra (some also say "Python 3.11+"). requirements.txt files using autolens[jax] keep working via the alias. Next: /start_workspace docs-migration pass.

Merge-order recommendation: PyAutoHeart#150 first (additive CI leg; callers pin lib-tests.yml@main, so the library PRs then exercise the no-jax leg), then the five library PRs bottom-up.

Post-release follow-up: after the first release with the promotion, bump the intra-family floors from >=2026.7.29.2 to that version (comments in each pyproject record this).

[Jammy2211]

All six library PRs merged; workspace PRs open behind pending-release

Merged (all suites + the new unittest-nojax leg green; main CI green on every repo): PyAutoHeart#150, PyAutoNerves#150, PyAutoFit#1503, PyAutoArray#450, PyAutoGalaxy#574, PyAutoLens#703.

The no-jax CI leg earned its keep on its first day — two real catches before anything reached main:

1. test_autolens/potential_correction/test_iterative_interferometer.py::test__solve_joint_optimization__identity_damping_finite reached the JAX-native sparse-operator path unmarked → @requires_jax added (94d8f54).
2. A genuine NumPy-path bug: BetaMessage/GammaMessage/NormalMessage misrouted NumPy scalars (np.int64/np.float64 are not int/float under NumPy 2) into the JAX import branch — np.generic added to the dispatch (PyAutoFit 19c679583). A jax-less user constructing messages from NumPy scalars would have crashed even before this promotion.

Cap widen reverted: jax 0.11.1 fails five autofit message-trace tests ('tuple' object has no attribute 'shape'); cap stays <0.11.0 (PyAutoNerves 848a254). Tracked in Mind draft/bug/autofit/jax_011_message_log_partition_tuple_shape.md — note the scalar-dispatch fix above covered test_beta, so the remaining 0.11 surface is likely just test_jax_trace.py.

Workspace docs PRs open, all pending-release (their prose claims default-JAX, true only from the promoted release): autolens_workspace#486, autogalaxy_workspace#212, autofit_workspace#139, HowToLens#71, HowToGalaxy#67. Notebooks regenerated via generate.py (unrelated drift excluded); markdown/ mirrors hand-patched with identical sentences (running generate_markdown.py executes every curated script — disproportionate for prose).

Remaining, release-blocked (nightly runs 02:00 UTC; today's merges arm it):

1. Bump intra-family floors >=2026.7.29.2 → the promoted release version in all five pyprojects.
2. Merge the five workspace PRs.
3. Later: make unittest-nojax a required status check once it has green history.