numpy==1.26.4
flax==0.8.3
jax[cpu]==0.4.28
optax==0.2.2

[gpu]
jax[cuda12]
