Docker: Upgrade Jax to 0.4.37

Dockerfile

@@ -95,6 +95,14 @@ ENV PIP_FIND_LINKS=https://storage.googleapis.com/jax-releases/libtpu_releases.h
 # Jax will fallback to CPU when run on a machine without TPU.
 RUN pip install .[core,tpu]
 RUN if [ -n "$EXTRAS" ]; then pip install .[$EXTRAS]; fi
+
+# V6E requires using jax 0.4.37 so temporarily override before AXLearn
+# upgrades to Jax 0.4.37.
+ENV DATE=20241201
+RUN pip install -U "jax[tpu]==0.4.37" "jax==0.4.37" "jaxlib==0.4.36" \
+    -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
+RUN pip install -U "libtpu-nightly==0.1.dev${DATE}" \
+    -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
 COPY . .
 
 ################################################################################