Skip to content

Adding support for jax and jaxoplanet#766

Draft
taylorbell57 wants to merge 80 commits intomainfrom
jaxoplanet
Draft

Adding support for jax and jaxoplanet#766
taylorbell57 wants to merge 80 commits intomainfrom
jaxoplanet

Conversation

@taylorbell57
Copy link
Collaborator

Because starry+PyMC3 are now so outdated that we can no longer maintain support for them, I am quickly working to add jax+jaxoplanet support.

This is just a draft PR to help make sure others know this feature is getting added soon, and I'll mark it as ready for review once it is.

This PR will resolve #170 and resolve #310

taylorbell57 and others added 22 commits October 22, 2024 12:06
Most models seem to work well now (polynomial, exponential ramp, and centroid have specifically been tested). The celerite2.jax GP model is currently failing, but it is also failing in the same way on the celerite2 docs pages, so I don't think that's an issue on my end.
@taylorbell57
Copy link
Collaborator Author

Reminding myself here that I can/should make progress on #591 in this PR

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request LC Fit

Projects

Status: No status
Status: In progress
Status: In progress

Development

Successfully merging this pull request may close these issues.

speed up GPs (try JAX, polychord, or implementing slow parts in C/C++) S5: Look into using jax based packages for both GP and nested sampling

1 participant