jax>=0.3.4
equinox
optax>=0.1.2
