Add support for PyMC v5 (and fix some recent Jax issues)

[vandalt]

Hi!

This adds support for PyMC > 5. Following the discussion in exoplanet-dev/exoplanet-core#89, I opted to replace the pymc4 versions with pymc directly and keep pymc3 only. Let me know if I should revert that.

Most of the changes just required replacing aesara with pytensor. A few deprecated functions and API changes also had to be replaced, but nothing major.

Also, while running tests, I noticed that both PyMC 4 and 5 gave errors when GaussianProcess.marginal() was called. The issue only happens when pytensor.config.compute_test_value != "off" ("debug mode"), which is the case only in the tests. After a bit of debugging, I think the bug comes from PyMC so I opened an issue (pymc-devs/pymc#6981) and a PR (pymc-devs/pymc#6982) to address it. I can confirm all tests are passing with a patched PyMC, but in the meantime regular code (without pytensor "debug mode") still works OK.

I also found some problems with deprecated JAX code, which prevented the PyMC+JAX support from working, so I fixed them in the same PR.

Some things I still have to do:

• Update Github workflows and noxfile.py
• Tutorials and documentation

Thanks!

EDIT: Added things left to do.

[dfm]

Thanks for this! In the tutorial, can you update the figure title to say "posterior psd using PyMC" instead of "PyMC3"? And maybe we could also remove the warning catching with PyMC - I think that was there for old Theano chatter.

[vandalt]

Thanks for the feedback!

can you update the figure title to say "posterior psd using PyMC" instead of "PyMC3"?

Done

And maybe we could also remove the warning catching with PyMC - I think that was there for old Theano chatter.

Also done. With PyMC v5 and PyTensor it only hides this one now:

/home/vandal/repos/astro/pymc/pymc/step_methods/hmc/quadpotential.py:629: UserWarning: QuadPotentialFullAdapt is an experimental feature
  warnings.warn("QuadPotentialFullAdapt is an experimental feature")

Personally I think it's not too verbose, and as a user I like seeing the same warnings as I get in my code when running a tutorial (to know they are expected and my setup is not broken, especially with something that involves compilation and that's new for many people, like PyMC). Let me know if you'd prefer to keep it hidden.

[dfm]

Thanks! I'm happy to have that warning in the docs (although it is a silly warning because that feature is several years old lol).

[dfm]

@vandalt — Take a look at the test logs when you get a chance. It looks like there are some issues with shape handling that will need some attention.

[vandalt]

Ah yes, this is what I was talking about here (in the original PR message):

Also, while running tests, I noticed that both PyMC 4 and 5 gave errors when GaussianProcess.marginal() was called. The issue only happens when pytensor.config.compute_test_value != "off" ("debug mode"), which is the case only in the tests. After a bit of debugging, I think the bug comes from PyMC so I opened an issue (pymc-devs/pymc#6981) and a PR (pymc-devs/pymc#6982) to address it. I can confirm all tests are passing with a patched PyMC, but in the meantime regular code (without pytensor "debug mode") still works OK.

The bug was indeed in PyMC and not in celerite2. The PyMC PR (pymc-devs/pymc#6982) is now merged so the issue is fixed when running the tests against the main Github branch. The options to make this pass are to wait for the next release, run the tests against the main branch, or disable the "compute_test_value flag.

[dfm]

Oh duh - sorry I forgot! I'd love to merge this, so perhaps we could disable the compute_test_value flag for those tests and then open an issue to remember to change it back after the next PyMC release. Thanks!

[vandalt]

Done @dfm! I added a TODO comment and can open an issue once the PR is merged (so I can link to the code lines). Thanks!