chex>=0.0.8
jax>=0.4.25
jaxlib>=0.4.25
