Skip to content

feat: promote jax_zero_contour to default dependency (PyAutoLens#702) #574

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Jammy2211 merged 1 commit into from
Aug 19, 2026
Merged

feat: promote jax_zero_contour to default dependency (PyAutoLens#702) #574

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions autogalaxy/plot/plot_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -357,8 +357,10 @@ def _critical_curves_method():
except ImportError:
logger.warning(
"critical_curves_method='zero_contour' requested, but "
"jax_zero_contour is not installed. Install autogalaxy[jax] "
"to use this method. Falling back to 'marching_squares'."
"jax_zero_contour is not installed. It is a default dependency "
"except on platforms without JAX wheels (e.g. Intel macOS); "
"install it with `pip install jax_zero_contour` to use this "
"method. Falling back to 'marching_squares'."
)
return "marching_squares"

Expand Down
6 changes: 4 additions & 2 deletions autogalaxy/util/plot_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -357,8 +357,10 @@ def _critical_curves_method():
except ImportError:
logger.warning(
"critical_curves_method='zero_contour' requested, but "
"jax_zero_contour is not installed. Install autogalaxy[jax] "
"to use this method. Falling back to 'marching_squares'."
"jax_zero_contour is not installed. It is a default dependency "
"except on platforms without JAX wheels (e.g. Intel macOS); "
"install it with `pip install jax_zero_contour` to use this "
"method. Falling back to 'marching_squares'."
)
return "marching_squares"

Expand Down
14 changes: 8 additions & 6 deletions docs/installation/conda.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,14 +37,16 @@ The latest version of **PyAutoGalaxy** is installed via pip as follows (the comm
caching issues impacting the installation):

```bash
pip install autogalaxy[jax] --no-cache-dir
pip install autogalaxy --no-cache-dir
```

The `[jax]` extra installs \[**JAX**\](<https://docs.jax.dev/en/latest/notebooks/thinking_in_jax.html>), which
**PyAutoGalaxy** uses for just-in-time compilation and GPU acceleration. **JAX is not installed by default** — to
install without it, use `pip install autogalaxy --no-cache-dir` instead. The extra installs CPU-only JAX; for GPU
support, follow the official \[JAX installation guide\](<https://jax.readthedocs.io/en/latest/installation.html>)
**before** installing.
This installs \[**JAX**\](<https://docs.jax.dev/en/latest/notebooks/thinking_in_jax.html>) by default, which
**PyAutoGalaxy** uses for just-in-time compilation and GPU acceleration (the older
`pip install autogalaxy[jax]` command still works and installs the same thing). The default install is CPU-only
JAX; for GPU support, follow the official
\[JAX installation guide\](<https://jax.readthedocs.io/en/latest/installation.html>) **before** installing.
On Intel (x86_64) macOS, where JAX publishes no wheels, the install automatically excludes JAX and runs on
the slower NumPy path — a warning is printed at import to make this clear.

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.
Expand Down
12 changes: 7 additions & 5 deletions docs/installation/pip.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,14 +24,16 @@ The latest version of **PyAutoGalaxy** is installed via pip as follows (specifyi
the installation has clean dependencies):

```bash
pip install autogalaxy[jax]
pip install autogalaxy
```

The `[jax]` extra installs \[**JAX**\](<https://docs.jax.dev/en/latest/notebooks/thinking_in_jax.html>), which
**PyAutoGalaxy** uses for just-in-time compilation and GPU acceleration. **JAX is not installed by default** — a
plain `pip install autogalaxy` gives a fully working install that runs on NumPy, without JAX acceleration. The
extra installs CPU-only JAX; for GPU support, follow the official
This installs \[**JAX**\](<https://docs.jax.dev/en/latest/notebooks/thinking_in_jax.html>) by default, which
**PyAutoGalaxy** uses for just-in-time compilation and GPU acceleration (the older
`pip install autogalaxy[jax]` command still works and installs the same thing). The default install is CPU-only
JAX; for GPU support, follow the official
\[JAX installation guide\](<https://jax.readthedocs.io/en/latest/installation.html>) **before** installing.
On Intel (x86_64) macOS, where JAX publishes no wheels, the install automatically excludes JAX and runs on
the slower NumPy path — a warning is printed at import to make this clear.

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.
Expand Down
25 changes: 10 additions & 15 deletions pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,15 +24,16 @@ classifiers = [
]
keywords = ["cli"]
dependencies = [
# Floors, not pins. Without them, pip backtracking the extras chain
# (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.
# PyAutoLens#687.
# Floors, not pins (PyAutoLens#687) — bump to the first release with JAX
# in the family's base dependencies once it exists (PyAutoLens#702), so
# backtracking cannot pair this autogalaxy with a jax-optional chain.
"autofit>=2026.7.29.2",
"autoarray>=2026.7.29.2",
"astropy>=5.0",
"nautilus-sampler==1.0.5"
"nautilus-sampler==1.0.5",
# Marker-gated like the jax deps in autonerves (jax_zero_contour depends
# on jax, which has no wheels on Intel macOS). PyAutoLens#702.
'jax_zero_contour>=2.0.0,<3.0.0; sys_platform != "darwin" or platform_machine == "arm64"'
]

[project.urls]
Expand All @@ -50,15 +51,9 @@ local_scheme = "no-local-version"


[project.optional-dependencies]
jax = [
# autofit[jax] rather than autonerves[jax] directly: it supplies everything
# autonerves[jax] does *plus* optax, which autofit's JAX-native gradient MAP
# searches (MultiStartAdam / MultiStartProdigy) import. Without it the chain
# autolens[jax] -> autogalaxy[jax] stopped at autonerves[jax], so `pip
# install autolens[jax]` left those searches raising ImportError.
"autofit[jax]>=2026.7.29.2",
"jax_zero_contour>=2.0.0,<3.0.0"
]
# JAX moved into the base dependencies (PyAutoLens#702). Kept as a declared
# no-op so `pip install autogalaxy[jax]` keeps resolving (PyAutoLens#687).
jax = []
optional = [
"autogalaxy[jax]",
"numba",
Expand Down
Loading