Repository navigation
feat: promote JAX from optional extra to default dependency #702
Description
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 basedependencies, 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_64wheels ≥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
- PyAutoNerves
pyproject.toml— move intodependencies: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) soautonerves[jax]stays valid; update the extras-chain comment block.
- PyAutoNerves
jax_wrapper— one-timewarnings.warn(or logger warning) when jax import fails: NumPy-only mode, performance significantly reduced, pointer to install docs. No behavioural change otherwise. - PyAutoFit
pyproject.toml— addoptax>=0.2.5todependencies; bumpautonervesfloor to the first promoted autonerves release;jax = []alias (keepoptaxthere or empty — alias only). - PyAutoArray
pyproject.toml— bumpautonervesfloor;jaxextra becomes alias.nufftax/pynufft/tfp-nightlystay inoptional(interferometer extras are NOT promoted). - PyAutoGalaxy
pyproject.toml— addjax_zero_contour>=2.0.0,<3.0.0todependencies; bumpautofit/autoarrayfloors; alias extra. - PyAutoLens
pyproject.toml— bumpautogalaxyfloor; alias extra. - All five repos — update
optionalextras 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. - PyAutoHeart
.github/workflows/lib-tests.yml— add aunittest-nojaxjob (Python 3.13 only): identical install, thenpip 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]). - Docs —
PyAutoLens/docs/installation/{pip,conda}.md(+ equivalent pages in the other repos' docs):pip install autolensis 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. - 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 widenPyAutoNerves/autonerves/**/jax_wrapper*.py— no-jax import warningPyAutoFit/pyproject.toml,PyAutoArray/pyproject.toml,PyAutoGalaxy/pyproject.toml,PyAutoLens/pyproject.toml— promotion + floors + alias extrasPyAutoHeart/.github/workflows/lib-tests.yml— no-jax CI legPyAutoLens/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 autolenstime; 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).
Activity
Jammy2211 commentedon Aug 19, 2026
Library PRs Created
All six PRs are open, every test suite green,
pending-releaselabel verified on the five library PRs: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.txtfiles usingautolens[jax]keep working via the alias. Next:/start_workspacedocs-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.2to that version (comments in each pyproject record this).test: mark solve_joint_optimization test requires_jax — caught by the…
Merge pull request #703 from PyAutoLabs/feature/jax-default-dependency
Jammy2211 commentedon Aug 19, 2026
All six library PRs merged; workspace PRs open behind pending-release
Merged (all suites + the new
unittest-nojaxleg green;mainCI 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:test_autolens/potential_correction/test_iterative_interferometer.py::test__solve_joint_optimization__identity_damping_finitereached the JAX-native sparse-operator path unmarked →@requires_jaxadded (94d8f54).BetaMessage/GammaMessage/NormalMessagemisrouted NumPy scalars (np.int64/np.float64are notint/floatunder NumPy 2) into the JAX import branch —np.genericadded 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 Minddraft/bug/autofit/jax_011_message_log_partition_tuple_shape.md— note the scalar-dispatch fix above coveredtest_beta, so the remaining 0.11 surface is likely justtest_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 viagenerate.py(unrelated drift excluded);markdown/mirrors hand-patched with identical sentences (runninggenerate_markdown.pyexecutes every curated script — disproportionate for prose).Remaining, release-blocked (nightly runs 02:00 UTC; today's merges arm it):
>=2026.7.29.2→ the promoted release version in all five pyprojects.unittest-nojaxa required status check once it has green history.feat: JAX default install, [jax] alias + install docs (PyAutoLabs#702)