jaxlib
jax
numpy
matplotlib
astropy
optax
numpyro
jaxopt
