From b0f3a3c8153310fc4ed06b9212d58887b725fb27 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 19 Aug 2026 14:10:39 -0400 Subject: [PATCH] feat: promote jax_zero_contour to default dependency (PyAutoLens#702) --- autogalaxy/plot/plot_utils.py | 6 ++++-- autogalaxy/util/plot_utils.py | 6 ++++-- docs/installation/conda.md | 14 ++++++++------ docs/installation/pip.md | 12 +++++++----- pyproject.toml | 25 ++++++++++--------------- 5 files changed, 33 insertions(+), 30 deletions(-) diff --git a/autogalaxy/plot/plot_utils.py b/autogalaxy/plot/plot_utils.py index 6f53dd27..a881a7be 100644 --- a/autogalaxy/plot/plot_utils.py +++ b/autogalaxy/plot/plot_utils.py @@ -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" diff --git a/autogalaxy/util/plot_utils.py b/autogalaxy/util/plot_utils.py index 6f53dd27..a881a7be 100644 --- a/autogalaxy/util/plot_utils.py +++ b/autogalaxy/util/plot_utils.py @@ -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" diff --git a/docs/installation/conda.md b/docs/installation/conda.md index b6bb0157..0e28dfb0 100644 --- a/docs/installation/conda.md +++ b/docs/installation/conda.md @@ -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**\](), 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\]() -**before** installing. +This installs \[**JAX**\]() 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\]() **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. diff --git a/docs/installation/pip.md b/docs/installation/pip.md index 498b8e32..3d7e7307 100644 --- a/docs/installation/pip.md +++ b/docs/installation/pip.md @@ -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**\](), 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**\]() 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\]() **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. diff --git a/pyproject.toml b/pyproject.toml index 40e0bcb9..2a1711bc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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] @@ -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",