Adding a pathways-tpu docker image

Author: lukebaumann

Dockerfile

@@ -92,6 +92,18 @@ RUN uv pip install --prerelease=allow .[core,tpu] && uv cache clean
 RUN if [ -n "$EXTRAS" ]; then uv pip install .[$EXTRAS] && uv cache clean; fi
 COPY . .
 
+################################################################################
+# Pathways-TPU container spec.                                                 #
+################################################################################
+
+FROM base AS pathways-tpu
+
+ARG EXTRAS=
+
+RUN uv pip install --prerelease=allow .[core,pathways-tpu] && uv cache clean
+RUN if [ -n "$EXTRAS" ]; then uv pip install .[$EXTRAS] && uv cache clean; fi
+COPY . .
+
 ################################################################################
 # GPU container spec.                                                          #
 ################################################################################