jax
flax
optax
