[fork CI check] Let the PyMC backend import without jax - #2
Closed
jdeast wants to merge 1 commit into
Closed
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.
Owner
Author
|
Superseded by exoplanet-dev#194 (this branch, opened upstream). Fork CI history retained here. |
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.
Fork-internal PR to exercise CI before opening upstream. Do not merge.