feat: JAX as a default dependency, marker-gated for Intel macOS (PyAutoLens#702)

autonerves/jax_wrapper.py

@@ -1,9 +1,25 @@
+import importlib.util
 import logging
 
 logger = logging.getLogger(__name__)
 
 import os
 
+if importlib.util.find_spec("jax") is None:
+    logger.warning(
+        """
+        JAX is not installed, so all computations will run on the pure NumPy
+        path. Performance is significantly reduced without JAX — model fits
+        that take minutes with JAX can take hours without it.
+
+        JAX is a default dependency of the PyAuto libraries; it is absent
+        either because this platform has no JAX wheels (e.g. Intel macOS) or
+        because it was uninstalled. On supported platforms, restore it with:
+
+            pip install jax
+        """
+    )
+
 xla_env = os.environ.get("XLA_FLAGS")
 
 xla_env_set = True

pyproject.toml

@@ -26,7 +26,18 @@ keywords = ["cli"]
 dependencies = [
     "typing-inspect>=0.4.0",
     "PyYAML>=6.0.1",
-    "numpy>=1.24.0,<3.0.0"
+    "numpy>=1.24.0,<3.0.0",
+    # JAX is a default dependency (PyAutoLens#702): CPU wheels exist for every
+    # supported platform except Intel macOS (jaxlib ships no macosx_x86_64
+    # wheels for >=0.7), so the markers let Intel Macs resolve to the
+    # NumPy-only path instead of failing at install. jax_wrapper warns loudly
+    # at import when JAX is absent.
+    # Cap stays <0.11: jax 0.11.1 breaks autofit's beta/gamma message
+    # log_partition under jit ('tuple' object has no attribute 'shape') —
+    # widen only together with that fix (tracked from PyAutoLens#702).
+    'jax>=0.7.0,<0.11.0; sys_platform != "darwin" or platform_machine == "arm64"',
+    'jaxlib>=0.7.0,<0.11.0; sys_platform != "darwin" or platform_machine == "arm64"',
+    'jaxnnls==1.0.1; sys_platform != "darwin" or platform_machine == "arm64"'
 ]
 
 [project.urls]
@@ -43,11 +54,11 @@ version_scheme = "post-release"
 local_scheme = "no-local-version"
 
 [project.optional-dependencies]
-jax = [
-    "jax>=0.7.0,<0.11.0",
-    "jaxlib>=0.7.0,<0.11.0",
-    "jaxnnls==1.0.1"
-]
+# JAX moved into the base dependencies (PyAutoLens#702). The extra is kept as
+# a declared no-op so `pip install autonerves[jax]` keeps resolving — removing
+# it would revive the pip history-walk trap ("version X does not provide the
+# extra 'jax'" is a warning, not an error; PyAutoLens#687) in reverse.
+jax = []
 optional = [
     "autonerves[jax]",
     "astropy>=5.0"