JAX brings automatic differentiation and the XLA compiler together through a NumPy-like API for high performance machine learning research on accelerators like GPUs and TPUs.
This is a curated list of awesome JAX libraries, projects, and other resources specifically coming from the field of astronomy/astrophysics. Contributions are welcome!
For a broader list of JAX resources, check out the awesome-jax repository.
- jax-euler-benchmarks - A simple JAX solver for the 2D Euler equations for benchmarking
- LINX - A fast, differentiable, and extensible public BBN code
- Papers: arXiv:2408.14538, arXiv:2408.14531
- flowMC - Normalizing-flow enhanced sampling package for probabilistic inference in Jax.
- Papers: arXiv:2211.06397, arXiv:2105.12603
- ripple - Differentiable Gravitational Waveforms with JAX
- Paper: arXiv:2302.05329
- jim - Gravitational-wave data analysis tools in Jax
- Papers: arXiv:2302.05333, arXiv:2404.11397
- fiesta - Fast inference of electromagnetic signals and transients with JAX
- jester - JAX-based EOS and TOV solver
- LINX - A fast, differentiable, and extensible public BBN code
- Papers: arXiv:2408.14538, arXiv:2408.14531