exoplanet-dev / celerite2 / 6739851294

78%

push

github

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

* import ShapedArray from core instead of abstract_arrays

* Use `xla_client.register_custom_call_target()` instead of deprecated `register_cpu_custom_call_target()`

Ref: https://github.com/google/jax/issues/17662

* pymc4 -> pymc and aesara -> pytensor

* Update from `DeviceArray` to `jax.Array`

Ref: https://jax.readthedocs.io/en/latest/jax_array_migration.html

* Implement `_supp_shape_from_params()` for `CeleriteNormalRV`

There use to be a default implementation of this but it has been removed.
Implemented using a reference parameter, as is done for most distributions.

* Replace deprecated `model.test_point` by `model.initial_point()`

* Use time as reference parameter for `CeleriteNormalRV`

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

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

* Move `pymc4.rst` to `pymc.rst` in docs

* Replace aesara_jax_fn by pytensor_jax_fn in test_pymc_ops.py

* Replace pymc4 by pymc in API docs

* Replace pymc4 by pymc in first.ipynb

* Replace pymc4 by pymc in workflows and noxfile

* PyMC3 -> PyMC in Getting Started (first.ipynb) tutorial PSD figure

* Uncomment pymc and pymc_jax Github workflows

Were commented because of errors with PyMC v4

* Remove warning filters when running PyMC sampling in Getting Started tutorial

* Upper limit on xarray version for PyMC3

* Comment `pytensor.config.compute_test_value = "raise"` for now

Requires a fix in PyMC. Introduced to main by PR #6982, commit 714b4a, will probably be in v5.9.2.

---------

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>

297 of 379 branches covered (0.0%)

Branch coverage included in aggregate %.

130 of 132 new or added lines in 6 files covered. (98.48%)

1970 of 2118 relevant lines covered (93.01%)

1.95 hits per line