jax
flax
optax
numpy
