jax
jaxlib
jaxtyping
numpy
equinox
optax
numpyro
