Update constraints to the latest stable

constraints_gpu.txt

@@ -1,65 +1,72 @@
-absl-py==1.4.0
+absl-py==2.1.0
+annotated-types==0.6.0
 aqtp==0.7.2
-array-record==0.5.0
+array_record==0.5.1
 astroid==3.1.0
 astunparse==1.6.3
 attrs==23.2.0
+black==24.3.0
+build==1.2.1
 cachetools==5.3.3
 certifi==2024.2.2
+cfgv==3.4.0
 charset-normalizer==3.3.2
-chex==0.1.85
+chex==0.1.86
 click==8.1.7
 cloud-accelerator-diagnostics==0.1.0
 cloud-tpu-diagnostics==0.1.5
 cloudpickle==3.0.0
 clu==0.0.12
 contextlib2==21.6.0
+dbus-python==1.2.18
 dill==0.3.8
+distlib==0.3.8
 dm-tree==0.1.8
 docstring_parser==0.16
 editdistance==0.8.1
 etils==1.7.0
-exceptiongroup==1.2.0
-flatbuffers==24.3.7
-flax==0.8.1
-fsspec==2024.2.0
-gast==0.4.0
-google-api-core==2.17.1
-google-auth==2.28.2
-google-auth-oauthlib==1.0.0
-google-cloud-aiplatform==1.47.0
+exceptiongroup==1.2.1
+filelock==3.13.4
+flatbuffers==24.3.25
+flax==0.8.2
+fsspec==2024.3.1
+gast==0.5.4
+google-api-core==2.18.0
+google-auth==2.29.0
+google-cloud-aiplatform==1.48.0
 google-cloud-appengine-logging==1.4.3
 google-cloud-audit-log==0.2.5
-google-cloud-bigquery==3.20.1
+google-cloud-bigquery==3.21.0
 google-cloud-core==2.4.1
 google-cloud-logging==3.10.0
 google-cloud-resource-manager==1.12.3
-google-cloud-storage==2.15.0
+google-cloud-storage==2.16.0
 google-crc32c==1.5.0
 google-jetstream==0.2.0
 google-pasta==0.2.0
 google-resumable-media==2.7.0
 googleapis-common-protos==1.63.0
-grain-nightly==0.0.6
+grain-nightly==0.0.7
 grpc-google-iam-v1==0.13.0
-grpcio==1.62.1
+grpcio==1.62.2
 grpcio-status==1.48.2
 gviz-api==1.10.0
-h5py==3.10.0
-idna==3.6
+h5py==3.11.0
+identify==2.5.36
+idna==3.7
 immutabledict==4.2.0
 importlab==0.8.1
-importlib_resources==6.3.0
+importlib_resources==6.4.0
 iniconfig==2.0.0
 isort==5.13.2
-jax==0.4.25
-jaxlib==0.4.25
+jax==0.4.26
+jaxlib==0.4.26+cuda12.cudnn89
 jaxtyping==0.2.28
 Jinja2==3.1.3
-keras==2.13.1
-libclang==16.0.6
-libcst==1.2.0
-Markdown==3.5.2
+keras==3.3.3
+libclang==18.1.1
+libcst==1.3.1
+Markdown==3.6
 markdown-it-py==3.0.0
 MarkupSafe==2.1.5
 mccabe==0.7.0
@@ -72,74 +79,77 @@ more-itertools==10.2.0
 msgpack==1.0.8
 msgspec==0.18.6
 mypy-extensions==1.0.0
+namex==0.0.8
 nest-asyncio==1.6.0
 networkx==3.1
 ninja==1.11.1.1
-numpy==1.24.3
-nvidia-cublas-cu12==12.4.2.65
-nvidia-cuda-cupti-cu12==12.4.99
-nvidia-cuda-nvcc-cu12==12.4.99
-nvidia-cuda-nvrtc-cu12==12.4.99
-nvidia-cuda-runtime-cu12==12.4.99
+nodeenv==1.8.0
+numpy==1.26.4
+nvidia-cublas-cu12==12.4.5.8
+nvidia-cuda-cupti-cu12==12.4.127
+nvidia-cuda-nvcc-cu12==12.4.131
+nvidia-cuda-nvrtc-cu12==12.4.127
+nvidia-cuda-runtime-cu12==12.4.127
 nvidia-cudnn-cu12==8.9.7.29
-nvidia-cufft-cu12==11.2.0.44
-nvidia-cusolver-cu12==11.6.0.99
-nvidia-cusparse-cu12==12.3.0.142
-nvidia-nccl-cu12==2.19.3
-nvidia-nvjitlink-cu12==12.4.99
-oauthlib==3.2.2
+nvidia-cufft-cu12==11.2.1.3
+nvidia-cusolver-cu12==11.6.1.9
+nvidia-cusparse-cu12==12.3.1.170
+nvidia-nccl-cu12==2.21.5
+nvidia-nvjitlink-cu12==12.4.127
 opt-einsum==3.3.0
-optax==0.2.1
-orbax-checkpoint==0.5.5
+optax==0.2.2
+optree==0.11.0
+orbax-checkpoint==0.5.10
 packaging==24.0
-pandas==2.2.1
-platformdirs==4.2.0
-pluggy==1.4.0
+pandas==2.2.2
+pathspec==0.12.1
+pip-tools==7.4.1
+platformdirs==4.2.1
+pluggy==1.5.0
 portpicker==1.6.0
 pre-commit==3.7.0
 promise==2.3
 proto-plus==1.23.0
 protobuf==3.20.3
 psutil==5.9.8
-pyasn1==0.5.1
-pyasn1-modules==0.3.0
+pyasn1==0.6.0
+pyasn1_modules==0.4.0
 pycnite==2023.10.11
-pydantic==1.10.14
+pydantic==2.7.1
+pydantic_core==2.18.2
 pydot==2.0.0
 pyglove==0.4.4
 Pygments==2.17.2
+PyGObject==3.42.1
 pyink==24.3.0
 pylint==3.1.0
 pyparsing==3.1.2
-pytest==8.1.1
+pyproject_hooks==1.0.0
+pytest==8.2.0
 python-dateutil==2.9.0.post0
-pytype==2024.3.11
+pytype==2024.4.11
 pytz==2024.1
-PyYAML==6.0.1
+PyYAML==5.4.1
 requests==2.31.0
-requests-oauthlib==1.4.0
 rich==13.7.1
 rsa==4.9
-scipy==1.12.0
+scipy==1.13.0
 sentencepiece==0.1.97
 seqio==0.0.19
-shapely==2.0.3
+shapely==2.0.4
 six==1.16.0
 tabulate==0.9.0
-tensorboard==2.13.0
+tensorboard==2.16.2
 tensorboard-data-server==0.7.2
 tensorboard_plugin_profile==2.15.1
 tensorboardX==2.6.2.2
-tensorflow==2.13.1
+tensorflow==2.16.1
 tensorflow-datasets==4.9.4
-tensorflow-estimator==2.13.0
-tensorflow-hub==0.16.1
 tensorflow-io-gcs-filesystem==0.36.0
-tensorflow-metadata==1.14.0
-tensorflow-text==2.13.0
-tensorstore==0.1.54
+tensorflow-metadata==1.15.0
+tensorflow-text==2.16.1
+tensorstore==0.1.58
 termcolor==2.4.0
-tf-keras==2.15.0
 tfds-nightly==4.9.2.dev202308090034
 toml==0.10.2
 tomli==2.0.1
@@ -148,10 +158,10 @@ toolz==0.12.1
 tqdm==4.66.2
 transformer-engine==1.5.0+297459b
 typeguard==2.13.3
-typing-inspect==0.9.0
-typing_extensions==4.5.0
+typing_extensions==4.11.0
 tzdata==2024.1
 urllib3==2.2.1
-Werkzeug==3.0.1
+virtualenv==20.26.0
+Werkzeug==3.0.2
 wrapt==1.16.0
-zipp==3.18.0
+zipp==3.18.1

docker_build_dependency_image.sh

@@ -57,7 +57,7 @@ if [[ -z ${LIBTPU_GCS_PATH+x} ]] ; then
   echo "Default LIBTPU_GCS_PATH=${LIBTPU_GCS_PATH}"
   if [[ ${DEVICE} == "gpu" ]]; then
     if [[ ${MODE} == "pinned" ]]; then
-      export BASEIMAGE=ghcr.io/nvidia/jax:base-2024-03-13
+      export BASEIMAGE=ghcr.io/nvidia/jax:base-2024-04-29
     else
       export BASEIMAGE=ghcr.io/nvidia/jax:base
     fi