flax
jax>=0.3.25
optax
tensorflow>=2.9.1
