Updating Pathways-TPU integration with AxLearn

Adding a pathways-tpu extra dependency kind
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.                                                          #
 ################################################################################

pyproject.toml

@@ -109,7 +109,12 @@ gcp = [
 tpu = [
     "axlearn[gcp]",
     "jax[tpu]==0.5.3",  # must be >=0.4.19 for compat with v5p.
-    "pathwaysutils==0.1.1",  # For JAX+Pathways single-controller accelerator coordinator.
+]
+# For Pathways-TPU single-controller training
+pathways-tpu = [
+    "axlearn[gcp]",
+    "jax==0.5.3",  # must be >=0.4.19 for compat with v5p.
+    "pathwaysutils==0.1.1",
 ]
 # Vertex AI tensorboard. TODO(markblee): Merge with `gcp`.
 vertexai_tensorboard = [