flax
optax
tensorflow-datasets