Fix MaxDiffusion outdated jax[tpu] version (GoogleCloudPlatform#1672)

* add rayjob dws

* add rayjob dws

* add rayjob dws

* add rayjob dws

* add rayjob dws

* add rayjob dws

* add rayjob dws

* add rayjob dws

* add rayjob dws

* add rayjob dws

* add rayjob dws

* add rayjob dws

* add rayjob dws

* add rayjob dws

* add rayjob dws

* update limit

* add file OSS header

* add file OSS header tag

* add file OSS header tag

* add file OSS header tag

* update resource

* add gpu

* add gpu

* add gpu

* add gpu

* add gpu

* add gpu

* update maxdiffusion dependency versions

* update maxdiffusion dependency versions

* update maxdiffusion dependency versions

* update maxdiffusion dependency versions

* update maxdiffusion dependency versions

* update maxdiffusion dependency versions

* update maxdiffusion dependency versions

* update jax[tpu] version

* update jax[tpu] version

Author: rick-c-goog

ai-ml/maxdiffusion-tpu/build/server/Dockerfile

@@ -6,7 +6,7 @@ RUN apt-get -y install git
 COPY requirements.txt ./
 RUN python -m pip install --upgrade pip
 RUN pip install -r requirements.txt
-RUN pip install git+https://github.com/google/maxdiffusion.git
+RUN pip install git+https://github.com/google/maxdiffusion.git@v1
 COPY main.py ./
 EXPOSE 8000
 ENTRYPOINT ["python", "main.py"]

ai-ml/maxdiffusion-tpu/build/server/requirements.txt

@@ -4,9 +4,10 @@
 uvicorn[standard]==0.30.6
 fastapi[all]==0.114.0
 pillow
-jax[tpu]==0.4.34.dev20241001 --pre
+jax[tpu]==0.4.36 --pre
 google-cloud-storage
 absl-py
+huggingface_hub==0.24.7
 transformers>=4.25.1
 datasets
 flax
@@ -23,5 +24,6 @@ parameterized
 Pillow
 pytest
 tensorflow==2.17.0
-tensorflow-datasets
+tensorflow-datasets>=4.9.6
 ruff>=0.6,<=0.6.4
+aqtp