diff --git a/autonerves/jax_wrapper.py b/autonerves/jax_wrapper.py index 0777513..72db1cd 100644 --- a/autonerves/jax_wrapper.py +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 14de217..f95f680 100644 --- a/pyproject.toml +++ b/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"