Skip to content

Fix JAX 0.8+ / PyMC v5 compat (supersedes #164)#165

Merged
dfm merged 28 commits intoexoplanet-dev:mainfrom
MilesCranmerBot:fix-jax
Feb 21, 2026
Merged

Fix JAX 0.8+ / PyMC v5 compat (supersedes #164)#165
dfm merged 28 commits intoexoplanet-dev:mainfrom
MilesCranmerBot:fix-jax

Conversation

@MilesCranmerBot
Copy link
Contributor

This PR is a re-opened/superseding version of MilesCranmer's original PR #164, with additional small fixes needed to get CI fully green on current dependency stacks.

Context / links:

What changed vs #164:

  • Build: add jax>=0.8.0 to [build-system].requires so the PEP517 build env can import jax during CMake configure and build/install the celerite2.jax.xla_ops extension (fixes ModuleNotFoundError: celerite2.jax.xla_ops).
  • JAX: remove dead import of jax.lib.xla_client (removed upstream in newer JAX).
  • JAX: update primitive implementation hook to use jax._src.dispatch.apply_primitive (and drop the pre-0.8 fallback since we require >=0.8).

All GitHub Actions checks are passing on the fork branch before opening this PR.

MilesCranmer and others added 28 commits November 23, 2025 01:12
Co-authored-by: Dan Foreman-Mackey <dfm@dfm.io>
Co-authored-by: Dan Foreman-Mackey <dfm@dfm.io>
Co-authored-by: Dan Foreman-Mackey <dfm@dfm.io>
Co-authored-by: Dan Foreman-Mackey <dfm@dfm.io>
Co-authored-by: Dan Foreman-Mackey <dfm@dfm.io>
Co-authored-by: Dan Foreman-Mackey <dfm@dfm.io>
Copy link
Member

@dfm dfm left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you!

@dfm dfm merged commit b3ea056 into exoplanet-dev:main Feb 21, 2026
6 of 7 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants