pin cudnn version

Author: priyakasimbeg

docker/Dockerfile

@@ -70,9 +70,7 @@ RUN if [ "$framework" = "jax" ] ; then \
         echo "Installing Jax GPU" \
         && cd /algorithmic-efficiency \
         && pip install -e '.[pytorch_cpu, full]' --extra-index-url https://download.pytorch.org/whl/cpu \
-        # Todo: remove temporary nightly install
         && pip install -e '.[jax_gpu]'; \
-        # &&  pip install -U --pre jax jaxlib "jax-cuda12-plugin[with-cuda]" jax-cuda12-pjrt -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/; \
     elif [ "$framework" = "pytorch" ] ; then \
         echo "Installing Pytorch GPU" \
         && cd /algorithmic-efficiency \

pyproject.toml

@@ -101,13 +101,14 @@ jax_core_deps = [
   "protobuf==4.25.5",
 ]
 jax_cpu = [
-  "jax==0.6.1",
+  "jax==0.7.0",
   "algoperf[jax_core_deps]",
 ]
 
 jax_gpu = [
-  "jax[cuda12]==0.6.1",
+  "jax[cuda12]==0.7.0",
   "algoperf[jax_core_deps]",
+  "nvidia-cudnn-cu12==9.10.2.21", # temporary workaround for https://github.com/jax-ml/jax/issues/30663
 ]
 
 pytorch_cpu = [