diff --git a/requirements.txt b/requirements.txt
index 7f56376..44c01f6 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -404,9 +404,9 @@ h5py==3.9.0 \
     --hash=sha256:eb7bdd5e601dd1739698af383be03f3dad0465fe67184ebd5afca770f50df9d6 \
     --hash=sha256:f68b41efd110ce9af1cbe6fa8af9f4dcbadace6db972d30828b911949e28fadd
     # via tensorflow
-idna==3.7 \
-    --hash=sha256:028ff3aadf0609c1fd278d8ea3089299412a7a8b9bd005dd08b9f8285bcb5cfc \
-    --hash=sha256:82fee1fc78add43492d3a1898bfa6d8a904cc97d8427f683ed8e798d07761aa0
+idna==3.15 \
+    --hash=sha256:048adeaf8c2d788c40fee287673ccaa74c24ffd8dcf09ffa555a2fbb59f10ac8 \
+    --hash=sha256:ca962446ea538f7092a95e057da437618e886f4d349216d2b1e294abfdb65fdc
     # via requests
 importlib-resources==6.1.0 \
     --hash=sha256:9d48dcccc213325e810fd723e7fbb45ccb39f6cf5c31f00cf2b965f5f10f3cb9 \
@@ -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