diff --git a/requirements.txt b/requirements.txt
index 7f56376..69c9578 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -421,19 +421,25 @@ jax==0.4.16 \
     #   flax
     #   optax
     #   orbax-checkpoint
-jaxlib==0.4.16 \
-    --hash=sha256:206da90ae8bf623f8b0cc6d0c01565bbbca5da596a9dac4ccf3fccb41bc1afa5 \
-    --hash=sha256:55e45e7af59909a8385b10e9d7072aab1dcb77eb65a16c63fe86b7344ee1543a \
-    --hash=sha256:82a8be6107db8590218e59078a770cdc3eb96153a394d6b38c61317c9e91d48d \
-    --hash=sha256:a5cdb37690e6bf31a29a3967f9f51c2ca22c35f81e525be15669e5f17784309e \
-    --hash=sha256:b9c1acbaf9bdeafe67c9efe7a7c85e7c7b6266185120fce1ce9156ddab468cd8 \
-    --hash=sha256:c790eab261ea11bec397b767c53e4dd25df418ab7938f669ab71c65fd3052a6a \
-    --hash=sha256:d249b28219a560fe0b145670c3d71869607d8af5495f974463642c514531e008 \
-    --hash=sha256:dff0870685127a89991a511039588b034379d4c92709bacd60c286b482d22eb6 \
-    --hash=sha256:e17c9a0e728432101b7cf12f3774420c48d74d66ce2fab56ad397261e89ba15b \
-    --hash=sha256:e19626b8b7ca286065a335be9d77494e70b051500045fe97b0009d0ed5c0c455 \
-    --hash=sha256:e402c7e652a802ed67bbbd7f3ffc1ef76dc9c92e269e273d1e9a3a9098e297ca \
-    --hash=sha256:e6cbd9e161bb340b2fef04430954a3783247ee45e4170cb768b74098f779d286
+jaxlib==0.6.1 \
+    --hash=sha256:02bac5153389f01616516a9fd1dcd6038d23ee50681dac14e4ddbc43ccb3133a \
+    --hash=sha256:11fcc4b1c741a1e0057f2ffa77d5a82bfe7ee97c3864ed88df67493e789b9173 \
+    --hash=sha256:2168217ec37bf951ca33377d3e0953178ba5cade95f194211d9ab2d53dcd2201 \
+    --hash=sha256:277cc7e9d657d0893a559261277b3eae916ad7fa73e300a629261fb537dca0f1 \
+    --hash=sha256:3301addee156f55d1f8079f80b314d89b80094740b7d64e5ec6e7ef2e1febbd7 \
+    --hash=sha256:5a90ee7c59b2c00773026fbf918269c7a8676a6a81a34a03af919f7d7bdce9a8 \
+    --hash=sha256:5e4f49113a527bcbac70c9e7074e95d8abfa35c3d67c2fed01f77a7abfd317aa \
+    --hash=sha256:76d6f65f3153ffb70e20a76b915d4431823cf70a786d86ba1b76a9c5bf66a0a4 \
+    --hash=sha256:7ae5815ada71b69532ce443a11160a3ed25c67e82a294a0d89af9d4d27429434 \
+    --hash=sha256:8106dc316eb440d07b9d4628a0c8e2acf76da5606742c9f5c33104aaa77b0ac2 \
+    --hash=sha256:acfe91eb44c29dbbd1f1f65f9bd66e1aef4483f57ad5e3d645129f3ec9ecde2a \
+    --hash=sha256:b12c8842b2dfc0770ca3785e183f7bed3fa1c2596c720591dbfbe29a05045108 \
+    --hash=sha256:b58c29fe747622b70946ea87823ad39202cc83da3d93a5293b432173b738a868 \
+    --hash=sha256:d039124468565bbf39363b1504c190e6719e6af89a7948dee256f1dee813bb94 \
+    --hash=sha256:d0c343c51b1052593edb603ddf58cf7f98812b2951ae6c45bd6e93e3e1f2f621 \
+    --hash=sha256:e14195c23eecd559a61c31027b4172e912e5a50f630320918ffdfae83090ca5a \
+    --hash=sha256:e734be70fe3e1fa2a31415362721189d974d10a66b0f5396c84585587d101b15 \
+    --hash=sha256:f4ca75d9d47a2e90099adfede0e9c926b83ef703d349b3289b8c88e861c09e5d
     # via
     #   chex
     #   optax
@@ -1203,9 +1209,9 @@ urllib3==2.0.7 \
     --hash=sha256:c97dfde1f7bd43a71c8d2a58e369e9b2bf692d1334ea9f9cae55add7d0dd0f84 \
     --hash=sha256:fdb6d215c776278489906c2f8916e6e7d4f5a9b602ccbcfdf7f016fc8da0596e
     # via requests
-werkzeug==3.0.1 \
-    --hash=sha256:507e811ecea72b18a404947aded4b3390e1db8f826b494d76550ef45bb3b1dcc \
-    --hash=sha256:90a285dc0e42ad56b34e696398b8122ee4c681833fb35b8334a095d82c56da10
+werkzeug==3.1.6 \
+    --hash=sha256:210c6bede5a420a913956b4791a7f4d6843a43b6fcee4dfa08a65e93007d0d25 \
+    --hash=sha256:7ddf3357bb9564e407607f988f683d72038551200c704012bb9a4c523d42f131
     # via tensorboard
 wheel==0.41.2 \
     --hash=sha256:0c5ac5ff2afb79ac23ab82bab027a0be7b5dbcf2e54dc50efe4bf507de1f7985 \