jax
numpy
equinox
