jax[cpu]
flax
numpy

