Add Python 3.13.0rc2 to JAX Docker images with CUDA 12.3 and CUDA 12.1.

PiperOrigin-RevId: 675594645
openxla · Sep 17, 2024 · 8ace4ee · 8ace4ee
1 parent 32295a0
commit 8ace4ee
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/third_party/tsl/tools/toolchains/remote_config/containers.bzl b/third_party/tsl/tools/toolchains/remote_config/containers.bzl
@@ -7,12 +7,12 @@ container_digests = {
     # JAX manylinux2014 configs.
     "cuda11.8-cudnn8.6-ubuntu20.04-manylinux2014-multipython": "sha256:45619e91f14faabddd79fe0cb1526df4c4ad92fc2e6ebdc725ea4419225429c3",
     "cuda12.1-cudnn8.9-ubuntu20.04-manylinux2014-multipython": "sha256:8c266e5b0acd203aed5e8871b63f68a39d8d23f6d882e619797e58b973f7fe63",
-    "cuda12.1-cudnn9.1-ubuntu20.04-manylinux2014-multipython": "sha256:7822b47867ecfc1f57df1cfadeaf091b72191d94cb722c271ed38809be7e7a61",
+    "cuda12.1-cudnn9.1-ubuntu20.04-manylinux2014-multipython": "sha256:7d948c3d2e3ab8867d600457b5666cc74c4206f08517791c95fc9a69b7cffefa",
     "cuda12.2-cudnn8.9-ubuntu20.04-manylinux2014-multipython": "sha256:9fefda035b4a12b24cd5bae56c7dbb9527a5fd06a41ced0a22ac86fe5ed26428",
     "cuda12.2-cudnn9.1-ubuntu20.04-manylinux2014-multipython": "sha256:0c78f3428cde36f041b758fc2f01d23d2f0dd72dec248f78667fb0c9d1f74cef",
     "cuda12.3-cudnn8.9-ubuntu20.04-manylinux2014-multipython": "sha256:6f9524a2ed7f75255dc4be3a0c5e3bda581385a1c13e2fa890bc17fa62da95b2",
     "cuda12.3-cudnn8.9-ubuntu22.04-manylinux2014-multipython": "sha256:97b219abb22994cf0530771d536f26fe301bacd328f0485c38af3847c2ee6b14",
-    "cuda12.3-cudnn9.1-ubuntu20.04-manylinux2014-multipython": "sha256:32cdd70157f0c4f82e0c0a0332d190227c1ac56b07715921f52972efdb0edfa1",
+    "cuda12.3-cudnn9.1-ubuntu20.04-manylinux2014-multipython": "sha256:e590303ea55a0990c26db4640161120ff6bc4124152c62155d397ba22d2ca850",
     # ROCM, probably not all of them still in use
     "rocm-ubuntu18.04-manylinux2010-multipython": "sha256:6e953a09b145df338bcb03e9e36f99b291140c29b72d0a048fb6c5905ccad5eb",
     "rocm-ubuntu20.04-manylinux2014-multipython": "sha256:906faec7765fe5dd067f2b092b5d5f220c1fedde725fb42c83d031b4d6f32204",

diff --git a/tools/toolchains/remote_config/containers.bzl b/tools/toolchains/remote_config/containers.bzl
@@ -7,12 +7,12 @@ container_digests = {
     # JAX manylinux2014 configs.
     "cuda11.8-cudnn8.6-ubuntu20.04-manylinux2014-multipython": "sha256:45619e91f14faabddd79fe0cb1526df4c4ad92fc2e6ebdc725ea4419225429c3",
     "cuda12.1-cudnn8.9-ubuntu20.04-manylinux2014-multipython": "sha256:8c266e5b0acd203aed5e8871b63f68a39d8d23f6d882e619797e58b973f7fe63",
-    "cuda12.1-cudnn9.1-ubuntu20.04-manylinux2014-multipython": "sha256:7822b47867ecfc1f57df1cfadeaf091b72191d94cb722c271ed38809be7e7a61",
+    "cuda12.1-cudnn9.1-ubuntu20.04-manylinux2014-multipython": "sha256:7d948c3d2e3ab8867d600457b5666cc74c4206f08517791c95fc9a69b7cffefa",
     "cuda12.2-cudnn8.9-ubuntu20.04-manylinux2014-multipython": "sha256:9fefda035b4a12b24cd5bae56c7dbb9527a5fd06a41ced0a22ac86fe5ed26428",
     "cuda12.2-cudnn9.1-ubuntu20.04-manylinux2014-multipython": "sha256:0c78f3428cde36f041b758fc2f01d23d2f0dd72dec248f78667fb0c9d1f74cef",
     "cuda12.3-cudnn8.9-ubuntu20.04-manylinux2014-multipython": "sha256:6f9524a2ed7f75255dc4be3a0c5e3bda581385a1c13e2fa890bc17fa62da95b2",
     "cuda12.3-cudnn8.9-ubuntu22.04-manylinux2014-multipython": "sha256:97b219abb22994cf0530771d536f26fe301bacd328f0485c38af3847c2ee6b14",
-    "cuda12.3-cudnn9.1-ubuntu20.04-manylinux2014-multipython": "sha256:32cdd70157f0c4f82e0c0a0332d190227c1ac56b07715921f52972efdb0edfa1",
+    "cuda12.3-cudnn9.1-ubuntu20.04-manylinux2014-multipython": "sha256:e590303ea55a0990c26db4640161120ff6bc4124152c62155d397ba22d2ca850",
     # ROCM, probably not all of them still in use
     "rocm-ubuntu18.04-manylinux2010-multipython": "sha256:6e953a09b145df338bcb03e9e36f99b291140c29b72d0a048fb6c5905ccad5eb",
     "rocm-ubuntu20.04-manylinux2014-multipython": "sha256:906faec7765fe5dd067f2b092b5d5f220c1fedde725fb42c83d031b4d6f32204",