Skip to content

Commit 6470e64

Browse files
Jammy2211Jammy2211
authored
Merge pull request #703 from PyAutoLabs/feature/jax-default-dependency
feat: JAX default install, [jax] alias + install docs (#702)
2 parents 6087581 + 94d8f54 commit 6470e64

4 files changed

Lines changed: 41 additions & 34 deletions

File tree

‎docs/installation/conda.md‎

Lines changed: 15 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -8,20 +8,21 @@
88

99
This acceleration is achieved through \[**JAX**\](<https://docs.jax.dev/en/latest/notebooks/thinking_in_jax.html>), which provides GPU and TPU support.
1010

11-
**JAX is not installed by default.** To install **PyAutoLens** with JAX, use the `jax` extra:
11+
**JAX is installed by default** — a plain `pip install autolens` includes it (the older
12+
`pip install autolens[jax]` command still works and installs the same thing).
1213

13-
```bash
14-
pip install autolens[jax] --no-cache-dir
15-
```
16-
17-
A plain `pip install autolens` gives a fully working install that runs on NumPy, but without any of the JAX
18-
acceleration described above.
19-
20-
The `[jax]` extra installs **CPU-only** JAX. To ensure GPU acceleration, it is recommended that you install JAX with
14+
The default install includes **CPU-only** JAX. To ensure GPU acceleration, it is recommended that you install JAX with
2115
GPU support **before** installing **PyAutoLens**, by following the official \[JAX installation guide\](<https://jax.readthedocs.io/en/latest/installation.html>).
2216

2317
If you install **PyAutoLens** without a proper GPU setup, a warning will be displayed.
2418

19+
:::{note}
20+
**Intel Macs**: JAX no longer publishes wheels for Intel (x86_64) macOS, so on these machines
21+
`pip install autolens` automatically installs without JAX and runs on the slower NumPy path — a
22+
warning is printed at import to make this clear. Every other supported platform (Windows, Linux,
23+
Apple-silicon Macs) gets JAX by default.
24+
:::
25+
2526
## Install
2627

2728
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
5758
caching issues impacting the installation):
5859

5960
```bash
60-
pip install autolens[jax] --no-cache-dir
61+
pip install autolens --no-cache-dir
6162
```
6263

63-
The `[jax]` extra is recommended, as it enables the JAX acceleration described above. To install without JAX,
64-
use `pip install autolens --no-cache-dir` instead.
64+
This includes JAX by default, enabling the acceleration described above. If you need an install without
65+
JAX on a platform where JAX wheels exist (e.g. a restricted environment), install normally and then run
66+
`pip uninstall jax jaxlib` — **PyAutoLens** detects the absence at import and falls back to the fully
67+
supported (but much slower) NumPy path.
6568

6669
If pip prints warnings about dependency version conflicts, these can usually be ignored — the instructions below
6770
will identify clearly if the installation is a success.

‎docs/installation/pip.md‎

Lines changed: 18 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -14,20 +14,21 @@ distribution" error. Upgrade Python to 3.12+ before installing.
1414

1515
This acceleration is achieved through \[**JAX**\](<https://docs.jax.dev/en/latest/notebooks/thinking_in_jax.html>), which provides GPU and TPU support.
1616

17-
**JAX is not installed by default.** To install **PyAutoLens** with JAX, use the `jax` extra:
17+
**JAX is installed by default** — a plain `pip install autolens` includes it (the older
18+
`pip install autolens[jax]` command still works and installs the same thing).
1819

19-
```bash
20-
pip install autolens[jax]
21-
```
22-
23-
A plain `pip install autolens` gives a fully working install that runs on NumPy, but without any of the JAX
24-
acceleration described above.
25-
26-
The `[jax]` extra installs **CPU-only** JAX. To ensure GPU acceleration, it is recommended that you install JAX with
20+
The default install includes **CPU-only** JAX. To ensure GPU acceleration, it is recommended that you install JAX with
2721
GPU support **before** installing **PyAutoLens**, by following the official \[JAX installation guide\](<https://jax.readthedocs.io/en/latest/installation.html>).
2822

2923
If you install **PyAutoLens** without a proper GPU setup, a warning will be displayed.
3024

25+
:::{note}
26+
**Intel Macs**: JAX no longer publishes wheels for Intel (x86_64) macOS, so on these machines
27+
`pip install autolens` automatically installs without JAX and runs on the slower NumPy path — a
28+
warning is printed at import to make this clear. Every other supported platform (Windows, Linux,
29+
Apple-silicon Macs) gets JAX by default.
30+
:::
31+
3132
## Install
3233

3334
We strongly recommend that you install **PyAutoLens** in a
@@ -40,20 +41,22 @@ We upgrade pip to ensure certain libraries install:
4041
pip install --upgrade pip
4142
```
4243

43-
The latest version of **PyAutoLens** is installed via pip as follows (specifying the version as shown below ensures
44-
the installation has clean dependencies):
44+
The latest version of **PyAutoLens** is installed via pip as follows:
4545

4646
```bash
47-
pip install autolens[jax]
47+
pip install autolens
4848
```
4949

50-
The `[jax]` extra is recommended, as it enables the JAX acceleration described above. To install without JAX,
51-
omit the extra:
50+
This includes JAX by default, enabling the acceleration described above. If you need an install without
51+
JAX on a platform where JAX wheels exist (e.g. a restricted environment), install normally and then remove it:
5252

5353
```bash
54-
pip install autolens
54+
pip uninstall jax jaxlib
5555
```
5656

57+
**PyAutoLens** detects the absence at import and falls back to the fully supported (but much slower)
58+
NumPy path.
59+
5760
If pip prints warnings about dependency version conflicts, these can usually be ignored — the instructions below
5861
will identify clearly if the installation is a success.
5962

‎pyproject.toml‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -24,12 +24,10 @@ classifiers = [
2424
]
2525
keywords = ["cli"]
2626
dependencies = [
27-
# Floor, not a pin. Without one, pip backtracking the extras chain
28-
# (autolens[jax] -> autogalaxy[jax] -> autofit[jax] -> autonerves[jax]) may
29-
# walk the release history to 2022: "version X does not provide the extra
30-
# 'jax'" is a pip *warning*, not an error, so a pre-extras release is a
31-
# legal solution. That is how `autolens[optional]` came to install autofit
32-
# 2026.4.30.582 and fail on `af.Latent` (#687).
27+
# Floor, not a pin (#687 — a floorless backtrack once installed autofit
28+
# 2026.4.30.582 and failed on `af.Latent`). Bump to the first release with
29+
# JAX in the family's base dependencies once it exists (#702), so
30+
# backtracking cannot pair this autolens with a jax-optional chain.
3331
"autogalaxy>=2026.7.29.2",
3432
"nautilus-sampler==1.0.5"
3533
]
@@ -49,7 +47,9 @@ local_scheme = "no-local-version"
4947

5048

5149
[project.optional-dependencies]
52-
jax = ["autogalaxy[jax]>=2026.7.29.2"]
50+
# JAX moved into the base dependencies (#702). Kept as a declared no-op so
51+
# `pip install autolens[jax]` keeps resolving (#687).
52+
jax = []
5353
coolest = ["coolest"]
5454
optional = [
5555
"autolens[jax]",

‎test_autolens/potential_correction/test_iterative_interferometer.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@ def test__solve_joint_optimization__finite_state_and_decreasing_cost(
6161
assert np.isfinite(dpsi_opt).all()
6262

6363

64+
@requires_jax
6465
def test__solve_joint_optimization__identity_damping_finite(interferometer_7):
6566
dataset = interferometer_7.apply_sparse_operator()
6667
fit = iter_fit_from(dataset, damping="identity", max_consecutive_rejections=3)

0 commit comments

Comments
 (0)