flax
jax
jaxlib
optax
