AttributeError: jax.lib.xla_client has no attribute "ops" when importing exoplanet_core.pymc

[fcotizelati]

I’m seeing an import-time crash with current JAX. The error originates in exoplanet_core.pymc where it accesses jax.lib.xla_client.ops, which no longer exists in recent JAX.

Environment:

• OS: macOS (Apple Silicon)
• Python: 3.14
• exoplanet-core: 0.3.1
• jax: 0.7.2
• jaxlib: 0.7.2
• pymc: 5.27.0
• pytensor: 2.36.2
• thejoker: 1.3.1 (error triggered via TheJoker.setup_mcmc)

Steps to reproduce:

1. Install current versions of exoplanet-core, jax, jaxlib
2. Run: from exoplanet_core.pymc import ops
3. Actual result:

Traceback (most recent call last):
  File "<python-input-5>", line 1, in <module>
    from exoplanet_core.pymc import ops
  File "/opt/homebrew/lib/python3.14/site-packages/exoplanet_core/pymc/__init__.py", line 28, in <module>
    from exoplanet_core.pymc import jax_support  # noqa
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/lib/python3.14/site-packages/exoplanet_core/pymc/jax_support.py", line 3, in <module>
    from exoplanet_core.jax import ops as jax_ops
  File "/opt/homebrew/lib/python3.14/site-packages/exoplanet_core/jax/__init__.py", line 23, in <module>
    from exoplanet_core.jax import ops  # noqa isort:skip
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/lib/python3.14/site-packages/exoplanet_core/jax/ops.py", line 16, in <module>
    xops = xla_client.ops
           ^^^^^^^^^^^^^^
  File "/opt/homebrew/lib/python3.14/site-packages/jax/_src/deprecations.py", line 57, in getattr
    raise AttributeError(f"module {module!r} has no attribute {name!r}")
AttributeError: module 'jax.lib.xla_client' has no attribute 'ops'

Expected result:
Import succeeds (or exoplanet-core pins a compatible JAX/JAXLIB range).

Additional context:
This shows up when calling TheJoker.setup_mcmc(...) which imports exoplanet_core.pymc.ops via exoplanet_core/_keplerian_orbit.py.