For DeepSeek inference pin JAX/MaxText to specific commits

Author: RajeshThallam

inference/trillium/JetStream-Maxtext/DeepSeek-R1-671B/docker/Dockerfile

@@ -44,8 +44,11 @@ ENV XCLOUD_ENVIRONMENT=GCP
 # Install JetStream and MaxText
 
 RUN git clone https://github.com/AI-Hypercomputer/JetStream.git && \
-git clone https://github.com/AI-Hypercomputer/maxtext.git && \
-git clone https://github.com/google/aqt.git
+    cd JetStream && git checkout 29329e8 && cd .. && \
+    git clone https://github.com/AI-Hypercomputer/maxtext.git && \
+    cd maxtext && git checkout 4f63ce61 && cd .. && \
+    git clone https://github.com/google/aqt.git && \
+    cd aqt && git checkout 28504f5 && cd ..
 
 RUN cd /maxtext && bash setup.sh && pip install torch --index-url https://download.pytorch.org/whl/cpu
 
@@ -56,4 +59,4 @@ RUN cd /JetStream && pip install -e .
 RUN apt -y update && apt-get -y install python3-dev && apt-get -y install build-essential
 RUN cp -r /aqt/aqt/* /usr/local/lib/python3.10/dist-packages/aqt/
 
-ENTRYPOINT [ "/bin/bash" ]
+ENTRYPOINT [ "/bin/bash" ]