Skip to content

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

Open
MilesCranmerBot wants to merge 27 commits intoexoplanet-dev:mainfrom
MilesCranmerBot:fix-jax
Open

Fix JAX 0.8+ / PyMC v5 compat (supersedes #164)#165
MilesCranmerBot wants to merge 27 commits intoexoplanet-dev:mainfrom
MilesCranmerBot:fix-jax

Conversation

@MilesCranmerBot
Copy link

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 27 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>
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.

2 participants