einops>=0.8.1
flax>=0.10.6
jax[cuda12]>=0.6.2
