Update Jax and Haiku versions

PiperOrigin-RevId: 474776295
Change-Id: I8f5add84f45069f2c48ac905801a611cb298b319

alphafold/model/utils.py

@@ -34,7 +34,7 @@ def final_init(config):
 
 def batched_gather(params, indices, axis=0, batch_dims=0):
   """Implements a JAX equivalent of `tf.gather` with `axis` and `batch_dims`."""
-  take_fn = lambda p, i: jnp.take(p, i, axis=axis)
+  take_fn = lambda p, i: jnp.take(p, i, axis=axis, mode='clip')
   for _ in range(batch_dims):
     take_fn = jax.vmap(take_fn)
   return take_fn(params, indices)

docker/Dockerfile

@@ -66,8 +66,8 @@ RUN wget -q -P /app/alphafold/alphafold/common/ \
 RUN pip3 install --upgrade pip \
     && pip3 install -r /app/alphafold/requirements.txt \
     && pip3 install --upgrade \
-      jax==0.2.14 \
-      jaxlib==0.1.69+cuda$(cut -f1,2 -d. <<< ${CUDA} | sed 's/\.//g') \
+      jax==0.3.17 \
+      jaxlib==0.3.15+cuda11.cudnn805 \
       -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 
 # Apply OpenMM patch.

docker/requirements.txt

 # Dependencies necessary to execute run_docker.py
-absl-py==0.13.0
+absl-py==1.0.0
 docker==5.0.0

requirements.txt

-absl-py==0.13.0
+absl-py==1.0.0
 biopython==1.79
 chex==0.0.7
-dm-haiku==0.0.4
+dm-haiku==0.0.7
 dm-tree==0.1.6
 docker==5.0.0
 immutabledict==2.0.0
-jax==0.2.14
+jax==0.3.17
 ml-collections==0.1.0
-numpy==1.19.5
+numpy==1.21.6
 pandas==1.3.4
 protobuf==3.20.1
 scipy==1.7.0
-tensorflow-cpu==2.5.0
+tensorflow-cpu==2.9.0