jax
jaxlib
equinox
optax
numpyro
