chex>=0.0.6

[jax]
jax>=0.2.17

[jaxlib]
jaxlib>=0.1.69
