feat: promote optax to default dependency; update no-jax messages (PyAutoLens#702)

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

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

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

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,

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):

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**\](<https://docs.jax.dev/en/latest/notebooks/thinking_in_jax.html>) (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\](<https://jax.readthedocs.io/en/latest/installation.html>)
-**before** installing.
+This installs \[**JAX**\](<https://docs.jax.dev/en/latest/notebooks/thinking_in_jax.html>) (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\](<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.
 
 Next, clone the `autofit_workspace` (the line `--depth 1` clones only the most recent branch on
 the `autofit_workspace`, reducing the download size):

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**\](<https://docs.jax.dev/en/latest/notebooks/thinking_in_jax.html>) (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**\](<https://docs.jax.dev/en/latest/notebooks/thinking_in_jax.html>) (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\](<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 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).

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]",