jax
flax
optax
chex<1.0.0
jaxtyping
numpy
scipy
gpjax
