!pip install jax jaxlib flax optax tensorflow tensorflow_datasets dm-pix tqdm matplotlib

        __
        __