Updated Jax to 0.3.25 and Haiku to 0.0.9

docker/Dockerfile

@@ -70,8 +70,8 @@ RUN wget -q -P /app/alphafold/alphafold/common/ \
 RUN pip3 install --upgrade pip --no-cache-dir \
     && pip3 install -r /app/alphafold/requirements.txt --no-cache-dir \
     && pip3 install --upgrade --no-cache-dir \
-      jax==0.3.17 \
-      jaxlib==0.3.15+cuda11.cudnn805 \
+      jax==0.3.25 \
+      jaxlib==0.3.25+cuda11.cudnn805 \
       -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 
 # Apply OpenMM patch.

requirements.txt

 absl-py==1.0.0
 biopython==1.79
 chex==0.0.7
-dm-haiku==0.0.7
+dm-haiku==0.0.9
 dm-tree==0.1.6
 docker==5.0.0
 immutabledict==2.0.0
-jax==0.3.17
+jax==0.3.25
 ml-collections==0.1.0
 numpy==1.21.6
 pandas==1.3.4