jax
jaxlib
jaxtyping
equinox
optax
numpyro
