Upgrade JAX to 0.4.26

djherbis · web-flow · commit 81a408971ede · 2024-05-03T10:44:06.000-04:00
http://b/338581280
diff --git a/config.txt b/config.txt
@@ -7,6 +7,6 @@ TORCH_VERSION=2.2.1
 TORCHAUDIO_VERSION=2.2.1
 TORCHTEXT_VERSION=0.17.1
 TORCHVISION_VERSION=0.17.1
-JAX_VERSION=0.4.23
+JAX_VERSION=0.4.26
 CUDA_MAJOR_VERSION=12
 CUDA_MINOR_VERSION=1
