jax>=0.2.0
jaxlib>=0.1.0
typing_extensions>=3.7.0
