Add extra libraries back to TPU (minus elegy)

elegy was the only package that was still failing the install

Author: djherbis

tpu/Dockerfile

@@ -56,10 +56,8 @@ ADD patches/kaggle_datasets.py /root/.local/lib/${PYTHON_VERSION_PATH}/site-pack
 RUN pip install tensorflow_hub https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/tensorflow/tf-${TENSORFLOW_VERSION}/tensorflow-${TENSORFLOW_VERSION}-${PYTHON_WHEEL_VERSION}-${PYTHON_WHEEL_VERSION}-${TF_LINUX_WHEEL_VERSION}.whl tensorflow-addons tensorflow-probability tensorflow-io \
     torch==${TORCH_VERSION} https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-${TORCH_VERSION}+libtpu-${PYTHON_WHEEL_VERSION}-${PYTHON_WHEEL_VERSION}-${TORCH_LINUX_WHEEL_VERSION}.whl torchvision==${TORCHVISION_VERSION} torchtext==${TORCHTEXT_VERSION} torchaudio==${TORCHAUDIO_VERSION} \
     jax[tpu]==${JAX_VERSION} -f https://storage.googleapis.com/jax-releases/libtpu_releases.html trax flax optax git+https://github.com/deepmind/dm-haiku jraph distrax \
-    # numpy==1.23.5 \
-    papermill jupyterlab python-lsp-server[all] "jupyter-lsp==1.5.1"
-    # elegy
-    # pandas matplotlib opencv-python-headless librosa accelerate diffusers scikit-learn transformers
+    papermill jupyterlab python-lsp-server[all] "jupyter-lsp==1.5.1" \
+    pandas matplotlib opencv-python-headless librosa accelerate diffusers scikit-learn transformers
 
 # Tensorflow libtpu:
 RUN curl --output /lib/libtpu.so https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/${TF_LIBTPU_VERSION}/libtpu.so