JAX 0.8 and PyMC v5 compatibility by MilesCranmer · Pull Request #164 · exoplanet-dev/celerite2

MilesCranmer · 2025-11-23T01:56:25Z

This PR upgrades celerite2 to JAX 0.8.x. I updated the jinja templates and re-generated the cpp files. I also upgraded PyMC to v5 and ditched compatibility with PyMC3 since stuff wasn't working anyways and it will be easier to maintain.

I'm unfamiliar with a lot of the lower level JAX stuff, so I am not confident about some of this PR, especially the FFI stuff. A good look over by someone else would be helpful.

Also, PyMC v5 seemed to need this @jax_funcify.register(_CeleriteOp) thing but I am not 100% sure about this.

Paging @dfm for review.

for more information, see https://pre-commit.ci

dfm

Thank you - this looks great! Some small comments/questions inline.

CMakeLists.txt

dfm · 2025-11-28T18:12:23Z

c++/include/celerite2/forward.hpp

  typedef typename LowRank::Scalar Scalar;
  typedef typename Eigen::internal::plain_col_type<Coeffs>::type CoeffVector;
-  typedef typename Eigen::Matrix<Scalar, LowRank::ColsAtCompileTime, RightHandSide::ColsAtCompileTime> Inner;
+  typedef typename Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor> Inner;


From my experience, this change will dramatically impact performance because Eigen won't be able to generate properly vectorized code for small systems. It's really useful to compile for specific sizes! Why did you make this change?

Thanks, I see. I couldn't get it working initially but this change seemed to do it. I didn't know it would hurt performance though so I'll fix it now.

dfm · 2025-11-28T18:12:30Z

c++/include/celerite2/forward.hpp

  typedef typename LowRank::Scalar Scalar;
  typedef typename Eigen::internal::plain_col_type<Coeffs>::type CoeffVector;
-  typedef typename Eigen::Matrix<Scalar, LowRank::ColsAtCompileTime, RightHandSide::ColsAtCompileTime> Inner;
+  typedef typename Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor> Inner;


pyproject.toml

python/spec/templates/jax/xla_ops.cpp

python/celerite2/jax/ops.py

Co-authored-by: Dan Foreman-Mackey <dfm@dfm.io>

for more information, see https://pre-commit.ci

Co-authored-by: Dan Foreman-Mackey <dfm@dfm.io>

MilesCranmer · 2025-11-29T20:58:17Z

Really struggling to get it working with LowRank::ColsAtCompileTime, RightHandSide::ColsAtCompileTime

MilesCranmer · 2025-11-29T21:55:49Z

Sorry, turns out I won't have enough bandwidth to fix the PR for fixed sized matrices (seems hard), so I will have to leave things as-is for now. Feel free to take this PR as-is if you are okay with the speed hit to get things up to JAX 0.8, or we can just point people to this PR if they need it

dfm · 2025-11-30T00:33:50Z

No problem! I think it shouldn't be a big deal since it's just for the "general" cases that are only used for predictions. Log prob calculations should still be fast. I'll try and get this merged soon - thanks!!

MilesCranmer · 2025-11-30T00:49:54Z

Oops sorry I forgot to re-run generate.py

MilesCranmer · 2025-12-11T10:48:57Z

Ping @dfm let me know if anything else is left!

dfm · 2025-12-15T01:39:36Z

It looks like all the CI is failing for various reasons. Can you take a look at those? I'm not totally sure how to have them automatically run for you, but I'll try to be faster to press the button, and you could plausibly run them on your own fork by temporarily adding:

- fix-jax

here:

celerite2/.github/workflows/python.yml

Line 5 in e7974e4

- main

It also looks like we'll need to update the Python version that we're using on ReadTheDocs. I think that should be a simple as just bumping it here:

celerite2/.readthedocs.yaml

Line 11 in e7974e4

python: "3.10"

Do you mind doing that too?

MilesCranmer · 2025-12-17T13:50:59Z

I'm not totally sure how to have them automatically run for you

I think if you put me as collaborator status in the org it might do this? I think it's just a user trust scopes thing. (No need to give me merge rights though)

dfm · 2026-01-03T19:59:39Z

Good idea! I've invited you - let me know if that works (or doesn't).

MilesCranmer · 2026-02-04T20:34:34Z

@MilesCranmerBot can you make a new PR based on this one and try to get the tests working again?

MilesCranmer · 2026-02-04T20:35:40Z

Actually ugh it might need to make the PR to my account's fork so I can get the CI tested via this PR

* fix: jax 0.8 compatibility * fix: get pymc working on v5 * deps: lower to python 3.10 * fix: JIT compat of `_do_compute` * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * deps: require python 3.11 for jaxlib compat * ci: skip bad nox arg * ci: fix nox target version * deps: fix missing jaxlib install * ci: fix missing jax install for pymc tests * fix: generator will no longer add lines of spaces Co-authored-by: Dan Foreman-Mackey <dfm@dfm.io> * fix: re-generated xla_ops without api.h Co-authored-by: Dan Foreman-Mackey <dfm@dfm.io> * refactor: clean up error handling Co-authored-by: Dan Foreman-Mackey <dfm@dfm.io> * fix: `Primitive` import Co-authored-by: Dan Foreman-Mackey <dfm@dfm.io> * build: update cmake per review Co-authored-by: Dan Foreman-Mackey <dfm@dfm.io> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * refactor: put axis into template Co-authored-by: Dan Foreman-Mackey <dfm@dfm.io> * fix: add back nrhs==1 branch * style: remove extra padding * deps: remove upper bound * chore: re-run generator * Fix JAX 0.8 primitive impl path and PyMC JAX registration * fix(jax): avoid removed jax.lib.xla_client import * build: require jax in PEP517 build env for JAX extension * refactor(jax): drop pre-0.8 apply_primitive fallback * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * build: avoid requiring jax on Python<3.11 for docs * Upgrade RTDs Python version from 3.10 to 3.11 --------- Co-authored-by: MilesCranmer <miles.cranmer@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Dan Foreman-Mackey <dfm@dfm.io> Co-authored-by: Dan Foreman-Mackey <foreman.mackey@gmail.com>

MilesCranmer and others added 5 commits November 23, 2025 01:12

fix: jax 0.8 compatibility

d6bf4e8

fix: get pymc working on v5

314ba3e

deps: lower to python 3.10

f13d3a4

fix: JIT compat of _do_compute

5790e09

[pre-commit.ci] auto fixes from pre-commit.com hooks

b14435b

for more information, see https://pre-commit.ci

MilesCranmer mentioned this pull request Nov 23, 2025

celerite2 fails with latest numpyro #157

Open

MilesCranmer added 5 commits November 23, 2025 02:12

deps: require python 3.11 for jaxlib compat

8b36a42

ci: skip bad nox arg

461fb35

ci: fix nox target version

9957276

deps: fix missing jaxlib install

7d5d97d

ci: fix missing jax install for pymc tests

272344d

dfm requested changes Nov 28, 2025

View reviewed changes

MilesCranmer and others added 10 commits November 29, 2025 17:19

fix: generator will no longer add lines of spaces

341c2f5

Co-authored-by: Dan Foreman-Mackey <dfm@dfm.io>

fix: re-generated xla_ops without api.h

ffaff43

Co-authored-by: Dan Foreman-Mackey <dfm@dfm.io>

refactor: clean up error handling

f6fc718

Co-authored-by: Dan Foreman-Mackey <dfm@dfm.io>

fix: Primitive import

2644c1f

Co-authored-by: Dan Foreman-Mackey <dfm@dfm.io>

build: update cmake per review

4fa899d

Co-authored-by: Dan Foreman-Mackey <dfm@dfm.io>

[pre-commit.ci] auto fixes from pre-commit.com hooks

f527c78

for more information, see https://pre-commit.ci

refactor: put axis into template

7679b85

Co-authored-by: Dan Foreman-Mackey <dfm@dfm.io>

fix: add back nrhs==1 branch

faa9b95

style: remove extra padding

63dd782

deps: remove upper bound

09d3ac4

chore: re-run generator

7d71090

dfm approved these changes Dec 15, 2025

View reviewed changes

This was referenced Feb 11, 2026

WIP: JAX 0.8 + PyMC v5 compat (from MilesCranmer PR #164) MilesCranmerBot/celerite2#1

Open

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

Merged

Comments

Conversation

MilesCranmer commented Nov 23, 2025

Uh oh!

dfm left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

dfm Nov 28, 2025

Choose a reason for hiding this comment

Uh oh!

MilesCranmer Nov 29, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dfm Nov 28, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

MilesCranmer commented Nov 29, 2025

Uh oh!

MilesCranmer commented Nov 29, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

dfm commented Nov 30, 2025

Uh oh!

MilesCranmer commented Nov 30, 2025

Uh oh!

MilesCranmer commented Dec 11, 2025

Uh oh!

dfm commented Dec 15, 2025

Uh oh!

MilesCranmer commented Dec 17, 2025

Uh oh!

dfm commented Jan 3, 2026

Uh oh!

MilesCranmer commented Feb 4, 2026

Uh oh!

MilesCranmer commented Feb 4, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

MilesCranmer Nov 29, 2025 •

edited

Loading

MilesCranmer commented Nov 29, 2025 •

edited

Loading