Let the PyMC backend import without jax - #194
Open
jdeast wants to merge 1 commit into
Open
Conversation
celerite2/pymc/ops.py imports pytensor.link.jax.dispatch at module scope,
which imports jax. That makes jax a hard requirement of the PyMC backend at
import time:
>>> from celerite2.pymc import terms
File "celerite2/pymc/ops.py", line 22, in <module>
from pytensor.link.jax.dispatch import jax_funcify
ModuleNotFoundError: No module named 'jax'
The only thing needed from that module is jax_funcify, and it is used
solely to register the JAX conversion -- which is useful exactly when jax
is installed. The conversion body already imports celerite2.jax.ops lazily
inside the function, so this is the last eager jax reference on the PyMC
path.
Guarding the import and registering conditionally leaves the JAX behaviour
identical where jax is present, and lets the PyMC and numpy backends work
where it is not. That matters on platforms jaxlib does not ship wheels for
-- macOS x86_64 stops at jaxlib 0.4.38, and jaxlib has no sdist -- where
installing jax at all is not an option.
Verified both ways against an installed celerite2 with this file patched
in: with jax importable, _CeleriteOp is still in jax_funcify.registry; with
sys.modules['jax'] poisoned to simulate its absence, 'from celerite2.pymc
import terms' succeeds.
jdeast
added a commit
to jdeast/EXOZIPPy
that referenced
this pull request
Aug 14, 2026
Moves the Intel-Mac caps out of the CI-only override script and into pyproject.toml as environment markers, so `poetry install` and `pip install exozippy` work on macOS x86_64. Everything here was measured on macos-15-intel across seven CI runs; see .github/workflows/intel-mac.yml. The pins: jax / jaxlib / numpyro / blackjax dropped on darwin+x86_64 numba >=0.62.1,<0.63 there (0.65.1 elsewhere) numpy <2.4 there, following numba 0.62.1 Dropping the jax family rather than pinning it to 0.4.38 is measured, not assumed, and it is the counterintuitive part: installing the newest jax that platform CAN hold is strictly WORSE than installing none. exoplanet-core reaches jax.ffi (public only since jax 0.5.0) and raises AttributeError, which escapes its own `except ImportError` guard, so `import exozippy` dies in components/orbit/orbit.py. With jax absent that guard catches a real ImportError and everything works. Upstream fix in flight: exoplanet-dev/exoplanet-core#146. numba is capped because PyTensor 3 makes it a hard dependency and its last macOS x86_64 wheels are 0.62.1 / llvmlite 0.45.1; past that the sdist wants a matching LLVM. PyTensor accepts numba>=0.58, so this stays inside its range. numpy follows from numba 0.62.1's own numpy<2.4. Verified: every dependency parses as PEP 508 and exactly ONE variant of each split requirement is selected on linux x86_64/aarch64, macOS arm64, macOS x86_64 and windows AMD64 -- a marker pair that both-selects or neither-selects would not show up until someone installed. poetry.lock regenerated; it now carries both numba entries with complementary markers. .github/scripts/intel_mac_deps.py is DELETED, as its docstring promised. With the markers committed, the workflow installs with a plain `pip install ".[gui]"` -- which is the point, since that is now exactly what a user gets. Keeping the generator would have left a second copy of the pin set to drift out of sync. Two gaps remain, both documented rather than papered over, and both with upstream fixes open: * `gp:` is unavailable -- celerite2's PyMC backend imports jax at module scope (exoplanet-dev/celerite2#194). The 10 kernel-building tests skip. * numpyro/blackjax are unavailable. nuts, ptde and nutpie all work. And one manual step: celerite2 0.3.3 must be pre-built with --no-build-isolation and BUILD_JAX=OFF, because its [build-system] requires names jax==0.8.0 and build isolation resolves that where our pins cannot reach (exoplanet-dev/celerite2#193). MACOS_INTEL_INSTALL.md is the runbook, linked from README's supported-platforms section, and says which upstream release deletes that step. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
celerite2/pymc/ops.pyimportspytensor.link.jax.dispatchat module scope, which imports jax. That makes jax a hard requirement of the PyMC backend at import time:The only thing needed from that module is
jax_funcify, and it is used solely to register the JAX conversion -- which is useful exactly when jax is installed. The conversion body already importscelerite2.jax.opslazily inside the function, so this is the last eager jax reference on the PyMC path.Guarding the import and registering conditionally leaves JAX behaviour identical where jax is present, and lets the PyMC and numpy backends work where it is not.
Why this matters
On platforms jaxlib publishes no wheels for, installing jax is not an option at all -- macOS x86_64 stops at jaxlib 0.4.38, and jaxlib has no sdist. Today that means
celerite2installs there but its PyMC backend cannot be imported, so GP modelling is unavailable even though nothing in it needs JAX.Verification
pymc_jax-- the one that would catch it if conditional registration had broken the JAX conversion where jax is present._CeleriteOpis still injax_funcify.registry;sys.modules['jax']poisoned to simulate absence,from celerite2.pymc import termssucceeds.Independent of #193, which fixes a separate build-time jax coupling; either can be merged without the other.
🤖 Generated with Claude Code