diff --git a/requirements.txt b/requirements.txt
index 7f56376..8bfeb0d 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
@@ -829,12 +835,13 @@ orbax-checkpoint==0.4.1 \
     --hash=sha256:309e6bbbffa0c90bb58121360d829e34607eaf28b70fc3e16469344307c24352 \
     --hash=sha256:9ed87757966e1d3bf1f20ac9ab39645a7a00263ee0a3a455ab73b0e635a25617
     # via flax
-packaging==23.2 \
-    --hash=sha256:048fb0e9405036518eaaf48a55953c750c11e1a1b68e0dd1a9d62ed0c092cfc5 \
-    --hash=sha256:8c491190033a9af7e1d931d0b5dacc2ef47509b34dd0de67ed209b5203fc88c7
+packaging==26.0 \
+    --hash=sha256:00243ae351a257117b6a241061796684b084ed1c516a08c48a3f7e147a9d80b4 \
+    --hash=sha256:b36f1fef9334a5588b4166f8bcd26a14e521f2b55e6b9de3aaa80d3ff7a37529
     # via
     #   matplotlib
     #   tensorflow
+    #   wheel
 pillow==10.3.0 \
     --hash=sha256:048ad577748b9fa4a99a0548c64f2cb8d672d5bf2e643a739ac8faff1164238c \
     --hash=sha256:048eeade4c33fdf7e08da40ef402e748df113fd0b4584e32c4af74fe78baaeb2 \
@@ -1207,9 +1214,9 @@ werkzeug==3.0.1 \
     --hash=sha256:507e811ecea72b18a404947aded4b3390e1db8f826b494d76550ef45bb3b1dcc \
     --hash=sha256:90a285dc0e42ad56b34e696398b8122ee4c681833fb35b8334a095d82c56da10
     # via tensorboard
-wheel==0.41.2 \
-    --hash=sha256:0c5ac5ff2afb79ac23ab82bab027a0be7b5dbcf2e54dc50efe4bf507de1f7985 \
-    --hash=sha256:75909db2664838d015e3d9139004ee16711748a52c8f336b52882266540215d8
+wheel==0.46.2 \
+    --hash=sha256:33ae60725d69eaa249bc1982e739943c23b34b58d51f1cb6253453773aca6e65 \
+    --hash=sha256:3d79e48fde9847618a5a181f3cc35764c349c752e2fe911e65fa17faab9809b0
     # via astunparse
 wrapt==1.14.1 \
     --hash=sha256:00b6d4ea20a906c0ca56d84f93065b398ab74b927a7a3dbd470f6fc503f95dc3 \