From dff3447a7445585d2b0448dc1f751aa411ae3f02 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 19 Aug 2026 14:07:45 -0400 Subject: [PATCH 1/2] feat: promote optax to default dependency; update no-jax messages (PyAutoLens#702) --- autofit/non_linear/analysis/analysis.py | 5 +++-- .../search/mle/multi_start_gradient/search.py | 6 ++++-- docs/installation/conda.md | 14 ++++++++------ docs/installation/pip.md | 12 +++++++----- pyproject.toml | 17 ++++++++++++----- 5 files changed, 34 insertions(+), 20 deletions(-) diff --git a/autofit/non_linear/analysis/analysis.py b/autofit/non_linear/analysis/analysis.py index a0527ded6..6a4df2af0 100644 --- a/autofit/non_linear/analysis/analysis.py +++ b/autofit/non_linear/analysis/analysis.py @@ -80,8 +80,9 @@ def __init__( "| Falling back to numpy. The fit will run, but JAX acceleration |\n" "| (typically 10-100x for large lens models) is unavailable. |\n" "| |\n" - "| To enable JAX, install on Python 3.11+ via your library's [jax] |\n" - "| extra, e.g.: pip install autolens[jax] |\n" + "| JAX is a default dependency; it is absent because this platform |\n" + "| has no JAX wheels (e.g. Intel macOS) or it was uninstalled. On |\n" + "| supported platforms, restore it with: pip install jax |\n" "+----------------------------------------------------------------------+", UserWarning, stacklevel=2, diff --git a/autofit/non_linear/search/mle/multi_start_gradient/search.py b/autofit/non_linear/search/mle/multi_start_gradient/search.py index 74eeefe9f..101727d45 100644 --- a/autofit/non_linear/search/mle/multi_start_gradient/search.py +++ b/autofit/non_linear/search/mle/multi_start_gradient/search.py @@ -651,8 +651,10 @@ def _fit( import optax.contrib # noqa: F401 — makes optax.contrib rules resolvable except ImportError as e: raise ImportError( - f"{type(self).__name__} requires the optional `jax` and `optax` " - "dependencies. Install them with `pip install autofit[jax] optax`." + f"{type(self).__name__} requires the `jax` and `optax` " + "dependencies. These are installed by default except on platforms " + "without JAX wheels (e.g. Intel macOS); install them with " + "`pip install jax optax`." ) from e if not getattr(analysis, "_use_jax", False): diff --git a/docs/installation/conda.md b/docs/installation/conda.md index 1cccd0462..b6edf3068 100644 --- a/docs/installation/conda.md +++ b/docs/installation/conda.md @@ -24,14 +24,16 @@ The latest version of **PyAutoFit** is installed via pip as follows (specifying the installation has clean dependencies): ```bash -pip install autofit[jax] +pip install autofit ``` -The `[jax]` extra installs \[**JAX**\]() (and -`optax`), which **PyAutoFit** uses for gradient-based searches and GPU acceleration. **JAX is not installed by -default** — to install without it, use `pip install autofit` instead. The extra installs CPU-only JAX; for GPU -support, follow the official \[JAX installation guide\]() -**before** installing. +This installs \[**JAX**\]() (and `optax`) by +default, which **PyAutoFit** uses for gradient-based searches and GPU acceleration (the older +`pip install autofit[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. Next, clone the `autofit_workspace` (the line `--depth 1` clones only the most recent branch on the `autofit_workspace`, reducing the download size): diff --git a/docs/installation/pip.md b/docs/installation/pip.md index df6410406..d24c94ba9 100644 --- a/docs/installation/pip.md +++ b/docs/installation/pip.md @@ -16,14 +16,16 @@ The latest version of **PyAutoFit** is installed via pip as follows (specifying the installation has clean dependencies): ```bash -pip install autofit[jax] +pip install autofit ``` -The `[jax]` extra installs \[**JAX**\]() (and -`optax`), which **PyAutoFit** uses for gradient-based searches and GPU acceleration. **JAX is not installed by -default** — a plain `pip install autofit` 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**\]() (and `optax`) by +default, which **PyAutoFit** uses for gradient-based searches and GPU acceleration (the older +`pip install autofit[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 this raises no errors **PyAutoFit** is installed! If there is an error check out the [troubleshooting section](https://pyautofit.readthedocs.io/en/latest/installation/troubleshooting.html). diff --git a/pyproject.toml b/pyproject.toml index 7539f046a..738ca5fd8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,11 +24,16 @@ classifiers = [ ] keywords = ["cli"] dependencies = [ - # Floor, not a pin. Without one, pip backtracking the extras chain - # (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. + # Floor, not a pin (PyAutoLens#687) — bump to the first release with JAX + # in autonerves' base dependencies once it exists (PyAutoLens#702), so + # backtracking cannot pair this autofit with a jax-optional autonerves. "autonerves>=2026.7.29.2", + # JAX-native gradient MAP searches (MultiStartAdam / MultiStartProdigy) + # import optax; jax itself comes from autonerves' base deps. Both are + # marker-gated for platforms with no jax wheels (Intel macOS) — optax + # depends on jax, so an unmarked optax would break resolution there. + # PyAutoLens#702. + 'optax>=0.2.5; sys_platform != "darwin" or platform_machine == "arm64"', "array_api_compat", "anesthetic>=2.9.0", "corner==2.2.2", @@ -68,7 +73,9 @@ local_scheme = "no-local-version" [project.optional-dependencies] -jax = ["autonerves[jax]>=2026.7.29.2", "optax>=0.2.5"] +# JAX moved into the base dependencies (PyAutoLens#702). Kept as a declared +# no-op so `pip install autofit[jax]` keeps resolving (PyAutoLens#687). +jax = [] mcp = ["mcp"] optional = [ "autofit[jax]", From 19c679583aeb8d6a599e746bb8755f2014774d9a Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 19 Aug 2026 14:51:49 -0400 Subject: [PATCH 2/2] =?UTF-8?q?fix:=20NumPy=20scalars=20misrouted=20messag?= =?UTF-8?q?e=20xp=20dispatch=20to=20the=20JAX=20branch=20=E2=80=94=20caugh?= =?UTF-8?q?t=20by=20the=20new=20no-jax=20CI=20leg=20(PyAutoLens#702)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit np.int64/np.float64 are not int/float subclasses under NumPy 2, so BetaMessage/GammaMessage/NormalMessage built from NumPy scalars imported jax (ModuleNotFoundError when absent). np.generic added to the NumPy-side isinstance check in all three. Co-Authored-By: Claude Fable 5 --- autofit/messages/beta.py | 2 +- autofit/messages/gamma.py | 2 +- autofit/messages/normal.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/autofit/messages/beta.py b/autofit/messages/beta.py index da1344c48..44e0fad19 100644 --- a/autofit/messages/beta.py +++ b/autofit/messages/beta.py @@ -146,7 +146,7 @@ def __init__( id_ Identifier for the message. Default is None. """ - if isinstance(alpha, (np.ndarray, float, int, list)): + if isinstance(alpha, (np.ndarray, np.generic, float, int, list)): xp = np else: import jax.numpy as jnp diff --git a/autofit/messages/gamma.py b/autofit/messages/gamma.py index 8414e0c38..c8cd0c8fe 100644 --- a/autofit/messages/gamma.py +++ b/autofit/messages/gamma.py @@ -29,7 +29,7 @@ def __init__( log_norm=0.0, id_=None ): - if isinstance(alpha, (np.ndarray, float, int, list)): + if isinstance(alpha, (np.ndarray, np.generic, float, int, list)): xp = np else: import jax.numpy as jnp diff --git a/autofit/messages/normal.py b/autofit/messages/normal.py index 1dbf1c927..2bb87c0de 100644 --- a/autofit/messages/normal.py +++ b/autofit/messages/normal.py @@ -106,7 +106,7 @@ def __init__( An optional unique identifier used to track the message in larger probabilistic graphs or models. """ - if isinstance(mean, (np.ndarray, float, int, list)): + if isinstance(mean, (np.ndarray, np.generic, float, int, list)): xp = np else: import jax.numpy as jnp