Generation of streams and subhalos using direct simulation or perturbation theory. Autodiffable everything, and GPU accelerated.

Unitful Quantities in JAX

Numerical quadrature with JAX

A multidimensional spline interpolation by Google/JAX

Nonlinear optimisation (root-finding, least squares, ...) in JAX+Equinox. https://docs.kidger.site/optimistix/

Interpolation and function approximation with JAX

Neural Network-Boosted Importance Nested Sampling for Bayesian Statistics

Efficient forward- and reverse-mode sparse Jacobians using Jax

Curated list of astronomy/astrophysics code packages using JAX

Tilting Disk Potential used in Nibauer et al. 2023: https://arxiv.org/abs/2312.09233