From 59c8a50e7a075bc4e5cd1c864dcd51cb14c648d2 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 19 Aug 2026 14:11:47 -0400 Subject: [PATCH 1/2] feat: JAX default install, [jax] alias + install docs (#702) --- docs/installation/conda.md | 27 +++++++++++++++------------ docs/installation/pip.md | 33 ++++++++++++++++++--------------- pyproject.toml | 14 +++++++------- 3 files changed, 40 insertions(+), 34 deletions(-) diff --git a/docs/installation/conda.md b/docs/installation/conda.md index 0512e68fc..2c59d406f 100644 --- a/docs/installation/conda.md +++ b/docs/installation/conda.md @@ -8,20 +8,21 @@ This acceleration is achieved through \[**JAX**\](), which provides GPU and TPU support. -**JAX is not installed by default.** To install **PyAutoLens** with JAX, use the `jax` extra: +**JAX is installed by default** — a plain `pip install autolens` includes it (the older +`pip install autolens[jax]` command still works and installs the same thing). -```bash -pip install autolens[jax] --no-cache-dir -``` - -A plain `pip install autolens` gives a fully working install that runs on NumPy, but without any of the JAX -acceleration described above. - -The `[jax]` extra installs **CPU-only** JAX. To ensure GPU acceleration, it is recommended that you install JAX with +The default install includes **CPU-only** JAX. To ensure GPU acceleration, it is recommended that you install JAX with GPU support **before** installing **PyAutoLens**, by following the official \[JAX installation guide\](). If you install **PyAutoLens** without a proper GPU setup, a warning will be displayed. +:::{note} +**Intel Macs**: JAX no longer publishes wheels for Intel (x86_64) macOS, so on these machines +`pip install autolens` automatically installs without JAX and runs on the slower NumPy path — a +warning is printed at import to make this clear. Every other supported platform (Windows, Linux, +Apple-silicon Macs) gets JAX by default. +::: + ## Install Installation via a conda environment circumvents compatibility issues when installing certain libraries. This guide @@ -57,11 +58,13 @@ The latest version of **PyAutoLens** is installed via pip as follows (the comman caching issues impacting the installation): ```bash -pip install autolens[jax] --no-cache-dir +pip install autolens --no-cache-dir ``` -The `[jax]` extra is recommended, as it enables the JAX acceleration described above. To install without JAX, -use `pip install autolens --no-cache-dir` instead. +This includes JAX by default, enabling the acceleration described above. If you need an install without +JAX on a platform where JAX wheels exist (e.g. a restricted environment), install normally and then run +`pip uninstall jax jaxlib` — **PyAutoLens** detects the absence at import and falls back to the fully +supported (but much slower) NumPy path. If pip prints warnings about dependency version conflicts, these can usually be ignored — the instructions below will identify clearly if the installation is a success. diff --git a/docs/installation/pip.md b/docs/installation/pip.md index e3dc31939..68bcc5a7a 100644 --- a/docs/installation/pip.md +++ b/docs/installation/pip.md @@ -14,20 +14,21 @@ distribution" error. Upgrade Python to 3.12+ before installing. This acceleration is achieved through \[**JAX**\](), which provides GPU and TPU support. -**JAX is not installed by default.** To install **PyAutoLens** with JAX, use the `jax` extra: +**JAX is installed by default** — a plain `pip install autolens` includes it (the older +`pip install autolens[jax]` command still works and installs the same thing). -```bash -pip install autolens[jax] -``` - -A plain `pip install autolens` gives a fully working install that runs on NumPy, but without any of the JAX -acceleration described above. - -The `[jax]` extra installs **CPU-only** JAX. To ensure GPU acceleration, it is recommended that you install JAX with +The default install includes **CPU-only** JAX. To ensure GPU acceleration, it is recommended that you install JAX with GPU support **before** installing **PyAutoLens**, by following the official \[JAX installation guide\](). If you install **PyAutoLens** without a proper GPU setup, a warning will be displayed. +:::{note} +**Intel Macs**: JAX no longer publishes wheels for Intel (x86_64) macOS, so on these machines +`pip install autolens` automatically installs without JAX and runs on the slower NumPy path — a +warning is printed at import to make this clear. Every other supported platform (Windows, Linux, +Apple-silicon Macs) gets JAX by default. +::: + ## Install We strongly recommend that you install **PyAutoLens** in a @@ -40,20 +41,22 @@ We upgrade pip to ensure certain libraries install: pip install --upgrade pip ``` -The latest version of **PyAutoLens** is installed via pip as follows (specifying the version as shown below ensures -the installation has clean dependencies): +The latest version of **PyAutoLens** is installed via pip as follows: ```bash -pip install autolens[jax] +pip install autolens ``` -The `[jax]` extra is recommended, as it enables the JAX acceleration described above. To install without JAX, -omit the extra: +This includes JAX by default, enabling the acceleration described above. If you need an install without +JAX on a platform where JAX wheels exist (e.g. a restricted environment), install normally and then remove it: ```bash -pip install autolens +pip uninstall jax jaxlib ``` +**PyAutoLens** detects the absence at import and falls back to the fully supported (but much slower) +NumPy path. + If pip prints warnings about dependency version conflicts, these can usually be ignored — the instructions below will identify clearly if the installation is a success. diff --git a/pyproject.toml b/pyproject.toml index fe893c2e0..2d2e841e0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,12 +24,10 @@ classifiers = [ ] keywords = ["cli"] dependencies = [ - # Floor, not a pin. Without one, pip backtracking the extras chain - # (autolens[jax] -> autogalaxy[jax] -> autofit[jax] -> autonerves[jax]) may - # walk the release history to 2022: "version X does not provide the extra - # 'jax'" is a pip *warning*, not an error, so a pre-extras release is a - # legal solution. That is how `autolens[optional]` came to install autofit - # 2026.4.30.582 and fail on `af.Latent` (#687). + # Floor, not a pin (#687 — a floorless backtrack once installed autofit + # 2026.4.30.582 and failed on `af.Latent`). Bump to the first release with + # JAX in the family's base dependencies once it exists (#702), so + # backtracking cannot pair this autolens with a jax-optional chain. "autogalaxy>=2026.7.29.2", "nautilus-sampler==1.0.5" ] @@ -49,7 +47,9 @@ local_scheme = "no-local-version" [project.optional-dependencies] -jax = ["autogalaxy[jax]>=2026.7.29.2"] +# JAX moved into the base dependencies (#702). Kept as a declared no-op so +# `pip install autolens[jax]` keeps resolving (#687). +jax = [] coolest = ["coolest"] optional = [ "autolens[jax]", From 94d8f54ba45ecb0c865b6c707cef0c7879cd429f Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 19 Aug 2026 14:41:33 -0400 Subject: [PATCH 2/2] =?UTF-8?q?test:=20mark=20solve=5Fjoint=5Foptimization?= =?UTF-8?q?=20test=20requires=5Fjax=20=E2=80=94=20caught=20by=20the=20new?= =?UTF-8?q?=20no-jax=20CI=20leg=20(#702)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Fable 5 --- .../potential_correction/test_iterative_interferometer.py | 1 + 1 file changed, 1 insertion(+) diff --git a/test_autolens/potential_correction/test_iterative_interferometer.py b/test_autolens/potential_correction/test_iterative_interferometer.py index 983d37f53..206d1adee 100644 --- a/test_autolens/potential_correction/test_iterative_interferometer.py +++ b/test_autolens/potential_correction/test_iterative_interferometer.py @@ -61,6 +61,7 @@ def test__solve_joint_optimization__finite_state_and_decreasing_cost( assert np.isfinite(dpsi_opt).all() +@requires_jax def test__solve_joint_optimization__identity_damping_finite(interferometer_7): dataset = interferometer_7.apply_sparse_operator() fit = iter_fit_from(dataset, damping="identity", max_consecutive_rejections=3)