jax[cuda13]==0.10.2
flax==0.12.8
optax==0.2.8
numpy==2.4.6
