[Docker] Add V7X requirements and update Docker to accept option to build using it (#916)

Signed-off-by: Jacob Platin <[email protected]>
Co-authored-by: Jacob Platin <[email protected]>

Author: jrplatin

docker/Dockerfile

@@ -3,6 +3,8 @@ ARG BASE_IMAGE="us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:night
 
 FROM $BASE_IMAGE
 
+ARG IS_FOR_V7X="false"
+
 # Remove existing versions of dependencies
 RUN pip uninstall -y torch torch_xla torchvision
 
@@ -39,4 +41,12 @@ RUN pip install -r requirements_benchmarking.txt
 COPY . .
 RUN pip install -e .
 
+# TODO (jacobplatin): remove when v7x is supported in JAX/Libtpu officially
+# NOTE: it's important that this is done after installing tpu_inference above,
+# so that the v7x-specific dependencies can override any existing ones.
+COPY requirements_v7x.txt .
+RUN if [ "$IS_FOR_V7X" = "true" ]; then \
+        pip install -r requirements_v7x.txt; \
+    fi
+
 CMD ["/bin/bash"]

requirements_v7x.txt

@@ -0,0 +1,8 @@
+# This file contains additional dependencies needed for TPU v7x support.
+# It is expected to be used in conjunction with the main requirements.txt file.
+--pre
+-i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/
+-f https://storage.googleapis.com/jax-releases/libtpu_releases.html
+jax==0.8.0.dev20251013
+jaxlib==0.8.0.dev20251013
+libtpu==0.0.25.dev20251012+nightly