jax[cpu]
equinox
