jax[cpu]
optax
equinox
addict
