Skip to content

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

Merged
Jammy2211 merged 2 commits into
mainfrom
feature/jax-default-dependency
Aug 19, 2026
Merged

feat: JAX as a default dependency, marker-gated for Intel macOS (PyAutoLens#702)#150
Jammy2211 merged 2 commits into
mainfrom
feature/jax-default-dependency

Conversation

@Jammy2211

@Jammy2211 Jammy2211 commented Aug 19, 2026

Copy link
Copy Markdown
Collaborator

Summary

Promotes JAX from the optional [jax] extra into the base dependencies (PyAutoLens#702): jax/jaxlib>=0.7,<0.11 (the planned cap widen to <0.12 is reverted: jax 0.11.1 breaks autofit's beta/gamma message log_partition under jit — tracked as a separate compat task) and jaxnnls==1.0.1 now install with a plain pip install autonerves. The three deps carry the environment marker sys_platform != "darwin" or platform_machine == "arm64" so Intel Macs — where jaxlib publishes no macosx_x86_64 wheels for >=0.7 — resolve cleanly to the NumPy-only path instead of failing at install. The [jax] extra is kept as a declared no-op alias so pip install autonerves[jax] keeps resolving (removing it would revive the pip history-walk trap of PyAutoLens#687 in reverse).

jax_wrapper now emits a one-time logger.warning at import when JAX is absent (NumPy-only mode, reduced performance, how to restore) — the NumPy fallback stays fully supported but is never silent.

API Changes

None — no Python API changes. Packaging: JAX moves from the [jax] extra into the default dependencies (marker-gated for Intel macOS); [jax] becomes a no-op alias.
See full details below.

Test Plan

  • pytest test_autonerves/ passes with JAX installed
  • pytest test_autonerves/ passes in a clean venv with JAX uninstalled (157/157), and the new import warning fires exactly once
  • pyproject parses; PEP 508 markers verified to include jax on linux/win/mac-arm64 and exclude it on Intel macOS
  • No-JAX CI leg (PyAutoHeart lib-tests.yml unittest-nojax) green once the Heart PR merges
Full API Changes (for automation & release notes)

Changed Behaviour

  • pip install autonerves now installs jax, jaxlib (>=0.7,<0.11) and jaxnnls by default on all platforms except Intel (x86_64) macOS.
  • autonerves.jax_wrapper logs a one-time warning at import when JAX is not installed.

Migration

  • Before: pip install autonerves[jax]
  • After: pip install autonerves (the [jax] form still works — no-op alias)

Part of the six-repo JAX-default-dependency change: PyAutoLens#702.

Generated by the PyAutoLabs agent workflow.

…ition (PyAutoLens#702)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@Jammy2211
Jammy2211 merged commit 4244856 into main Aug 19, 2026
5 checks passed
@Jammy2211
Jammy2211 deleted the feature/jax-default-dependency branch August 19, 2026 19:05
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

pending-release PR queued for the next release build

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant