diff --git a/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix b/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix
index dd45c995d18d..affa3461f26a 100644
--- a/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix
+++ b/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix
@@ -17,11 +17,11 @@ let
   platforms = {
     x86_64-linux = {
       name = "manylinux_2_27_x86_64";
-      hash = "sha256-ITOzr8YzB8AWrkmA8Ways94Zl1k2Pq1099VQ6q6VNGM=";
+      hash = "sha256-XsKOBhmNCUAxw3WQtOp1OujN5IeatLYLr0Dw3/tnMoY=";
     };
     aarch64-linux = {
       name = "manylinux_2_27_aarch64";
-      hash = "sha256-82JH1AlzsgrPPGnax8nHINqqQJaoCbPo/XApK+YXzes=";
+      hash = "sha256-5qpsWMVCDBvaItsxovbFwdCmbLcB7rcx/q5E/vGXcXw=";
     };
   };
   currentPlatform = platforms.${stdenv.hostPlatform.system};
diff --git a/pkgs/development/python-modules/jax-cuda12-plugin/default.nix b/pkgs/development/python-modules/jax-cuda12-plugin/default.nix
index 5bc26ccee746..a2cbc8d9b506 100644
--- a/pkgs/development/python-modules/jax-cuda12-plugin/default.nix
+++ b/pkgs/development/python-modules/jax-cuda12-plugin/default.nix
@@ -19,17 +19,17 @@ let
     x86_64-linux = {
       name = "manylinux_2_27_x86_64";
       hashes = {
-        cp312 = "sha256-rmekzEmzTFdXE6gxH3WX9nAOxaK2r728PbVVD7Aihzw=";
-        cp313 = "sha256-Ve1tvfTW9XxsewqtrrjjNzcaJSSaDjwe+e/LYHsLvGc=";
-        cp314 = "sha256-ltGYFX526tkZWHQkZvXHRiYY3SeDHcgjhIgOLoDr5OA=";
+        cp312 = "sha256-YdxnqDQa1qC7WWMTB5FFga0WBV9sVjWit+V9//t35GY=";
+        cp313 = "sha256-Vn3g7Q5wiq6vmIULqGmxVtkbIBRRlosGvu/DMdpJMp8=";
+        cp314 = "sha256-29SW/86v1w3lRyZciw/vbjsJU/WbQKYr3qJMZxUm4Qk=";
       };
     };
     aarch64-linux = {
       name = "manylinux_2_27_aarch64";
       hashes = {
-        cp312 = "sha256-rQ+7YFSFhdXw3y/hHrjm1BqYVidO2bw+vi0eHhw9AS8=";
-        cp313 = "sha256-KRqt6kpSxTo1+VfKWjwN1TFrXAgvy7UfFkt/INTh1lk=";
-        cp314 = "sha256-7H4G16RQgHln81CR1K9sCnJGG9x3xLIXGKUNg3fq6vs=";
+        cp312 = "sha256-qIHb70FjEPUC6sY3ZU+and5/5LMZ3lehg4dfiG3WKsc=";
+        cp313 = "sha256-4pwRBYRMHaWajYHApfF9lNKFJIF7QOHJITwEUAbZThc=";
+        cp314 = "sha256-5q/fhhyOuwAiYWYAXD2VaO4GM0npGnR+1gCyxEBb1Sc=";
       };
     };
   };
diff --git a/pkgs/development/python-modules/jax-cuda13-pjrt/default.nix b/pkgs/development/python-modules/jax-cuda13-pjrt/default.nix
index e5f9becb63e9..30b3b301e1b2 100644
--- a/pkgs/development/python-modules/jax-cuda13-pjrt/default.nix
+++ b/pkgs/development/python-modules/jax-cuda13-pjrt/default.nix
@@ -17,11 +17,11 @@ let
   platforms = {
     x86_64-linux = {
       name = "manylinux_2_27_x86_64";
-      hash = "sha256-5f9YnUMGSRO7dAJ2gtiWOcfqSPn4HCuub8dRF+kT+T4=";
+      hash = "sha256-WJm7YXftqG6lMq+5omeU8r3QSWb2oB+QKCZrFsK2z5A=";
     };
     aarch64-linux = {
       name = "manylinux_2_27_aarch64";
-      hash = "sha256-Y7VeUOXCWDol4Za6BrA0tj0rLYrUwvF5NWUsjTrS3S0=";
+      hash = "sha256-QHWNBC34wGmRtqGoAWtW6c77HtTGd8BiElw+txw9q5c=";
     };
   };
   currentPlatform = platforms.${stdenv.hostPlatform.system};
diff --git a/pkgs/development/python-modules/jax-cuda13-plugin/default.nix b/pkgs/development/python-modules/jax-cuda13-plugin/default.nix
index 995dbcb99200..59be239c33fc 100644
--- a/pkgs/development/python-modules/jax-cuda13-plugin/default.nix
+++ b/pkgs/development/python-modules/jax-cuda13-plugin/default.nix
@@ -19,17 +19,17 @@ let
     x86_64-linux = {
       name = "manylinux_2_27_x86_64";
       hashes = {
-        cp312 = "sha256-rTU+0YdyYQVg36KNXhVNntN3/SVyggpI9bVZEDe4eQk=";
-        cp313 = "sha256-l5aBOt4ZkL18neR1ubqmqidclzYSOQEb/DICKXZvLV4=";
-        cp314 = "sha256-hMDOpE60Ju825kzR3RRhmauiFWFyMSOQlpBHwmvM7qM=";
+        cp312 = "sha256-YCj68XnnqnzGQQDtPHx7GxE0b1qgXV6eM+j+YMUon6E=";
+        cp313 = "sha256-SBcsMa7srjEhwFYbVNm4otDHVUk1l5jc7/PluevPdN8=";
+        cp314 = "sha256-FrfbYs0i21END/Uv268dOxSQic/2oZPRldkinY6xQ3k=";
       };
     };
     aarch64-linux = {
       name = "manylinux_2_27_aarch64";
       hashes = {
-        cp312 = "sha256-AWY4p5REkWPKnN1ZvMvQ44AovyjnMhfLlryLFCpbjOI=";
-        cp313 = "sha256-TIhgusB4zeTetWV6o4l7A5lM/GYngGiUc1QPL+89Nzw=";
-        cp314 = "sha256-P0ZWPnmhGL359m+VmMeFrjxtIX74NPE4Qx6OdCZDlPM=";
+        cp312 = "sha256-CH8TMc1asn3NQwO61xKe+vfx3u2DDVZLjQXX7+yUbWE=";
+        cp313 = "sha256-hGeGn/1kvpd7DzwiuY/B+NYBl3Z8ZI9qLHv7/1ZIlK4=";
+        cp314 = "sha256-in5K+eDFVwsue9VCOva2hm26pxndlR0kF7lIBvfJFCs=";
       };
     };
   };
diff --git a/pkgs/development/python-modules/jax/default.nix b/pkgs/development/python-modules/jax/default.nix
index 83c3206c995f..21d858d6fd2b 100644
--- a/pkgs/development/python-modules/jax/default.nix
+++ b/pkgs/development/python-modules/jax/default.nix
@@ -44,7 +44,7 @@ let
 in
 buildPythonPackage (finalAttrs: {
   pname = "jax";
-  version = "0.11.2";
+  version = "0.12.0.1";
   pyproject = true;
   __structuredAttrs = true;
 
@@ -53,7 +53,7 @@ buildPythonPackage (finalAttrs: {
     repo = "jax";
     # google/jax contains tags for jax and jaxlib. Only use jax tags!
     tag = "jax-v${finalAttrs.version}";
-    hash = "sha256-lBfsbEa5Z6GmGVSgdDKoMF0R1TuU+e05QYP96qPNnxQ=";
+    hash = "sha256-um7ZO79sZCVFW0Ci4igjabBXildH698V1mCKpIBUVe8=";
   };
 
   build-system = [ setuptools ];
diff --git a/pkgs/development/python-modules/jaxlib/bin.nix b/pkgs/development/python-modules/jaxlib/bin.nix
index 596e8476705a..0377e9b5fd66 100644
--- a/pkgs/development/python-modules/jaxlib/bin.nix
+++ b/pkgs/development/python-modules/jaxlib/bin.nix
@@ -17,7 +17,7 @@
 }:
 
 let
-  version = "0.11.2";
+  version = "0.12.0.1";
   inherit (python) pythonVersion;
 
   # As of 2023-06-06, google/jax upstream is no longer publishing CPU-only wheels to their GCS bucket. Instead the
@@ -47,49 +47,49 @@ let
       "3.12-x86_64-linux" = getSrcFromPypi {
         platform = "manylinux_2_27_x86_64";
         dist = "cp312";
-        hash = "sha256-Yzih8JVvYNAlKUBKwzqAcAFaUmNc22u7FdDRmXVwd7s=";
+        hash = "sha256-eDGOmseckpcvgcdwkPVXvbSlpJghdRxu4/apxmqUiEI=";
       };
       "3.12-aarch64-linux" = getSrcFromPypi {
         platform = "manylinux_2_27_aarch64";
         dist = "cp312";
-        hash = "sha256-1JcAToDaqRuYJY1kWzdOjf9IC/AX0sXz4wsTcayPVLE=";
+        hash = "sha256-QIEiAcO81MsQS4j3H1j+xMvQwyupDlySs/qPQIGfZ3Q=";
       };
       "3.12-aarch64-darwin" = getSrcFromPypi {
         platform = "macosx_11_0_arm64";
         dist = "cp312";
-        hash = "sha256-7KgOdFt2IzaadrX2hWduwjOjOlXWn2D6flfL78/Uetw=";
+        hash = "sha256-2yzIWE19DRZik+8kdbcmmIoQ9vm1qOO443EkFeVXA4A=";
       };
 
       "3.13-x86_64-linux" = getSrcFromPypi {
         platform = "manylinux_2_27_x86_64";
         dist = "cp313";
-        hash = "sha256-XzysjXAwwfgKGCAl1xceo7EkUWKnEPbljSXuP+F0iqw=";
+        hash = "sha256-/cVpmbppC9kV9wKWEyaYSOqvhdIx7Iq8RmENO4QMGJw=";
       };
       "3.13-aarch64-linux" = getSrcFromPypi {
         platform = "manylinux_2_27_aarch64";
         dist = "cp313";
-        hash = "sha256-X504M8sr6l40a9ITi6dJkgF+CEHGPYkKAVehP5y/Q54=";
+        hash = "sha256-yN4crTgnExFPZX9AY4PtxT7p/aRnnbTisJ0d/v4cTkU=";
       };
       "3.13-aarch64-darwin" = getSrcFromPypi {
         platform = "macosx_11_0_arm64";
         dist = "cp313";
-        hash = "sha256-dmvZDieg/1O4e7nXDHNWXsgLmxFlMgzLj3LeRrhVEN8=";
+        hash = "sha256-i+FgsLhHvFcWjkMFcTdW9toNC7omcV5k+dL0bAp6a0I=";
       };
 
       "3.14-x86_64-linux" = getSrcFromPypi {
         platform = "manylinux_2_27_x86_64";
         dist = "cp314";
-        hash = "sha256-HzHYprH7Exgf8YB3V1BtEPbny6aycH9/gPc7hTyy/yI=";
+        hash = "sha256-tgeadmnWJOvULRVxM5Vm5iz+ciX0taILjhZ3PqmoNXo=";
       };
       "3.14-aarch64-linux" = getSrcFromPypi {
         platform = "manylinux_2_27_aarch64";
         dist = "cp314";
-        hash = "sha256-BhRo61rGtiEyFf0RArI3vLyT8kSAyIEkW3NhnQwqMh4=";
+        hash = "sha256-YJ3v+uo4jR7Z0cUQTY9gQXkfAHca3DBfcxa3QvUC/Ws=";
       };
       "3.14-aarch64-darwin" = getSrcFromPypi {
         platform = "macosx_11_0_arm64";
         dist = "cp314";
-        hash = "sha256-7a1rKhjWPcKWTMl6FtEm2tQWIe2sedZR4PbIeF2JE2k=";
+        hash = "sha256-k7jaGD8CNWA8RgtTDajNbrI+p/HGrcT4KkOmA1SX1GY=";
       };
     };
 in
diff --git a/pkgs/development/python-modules/jaxopt/default.nix b/pkgs/development/python-modules/jaxopt/default.nix
index 2ea513c2a09e..6b605bc1c989 100644
--- a/pkgs/development/python-modules/jaxopt/default.nix
+++ b/pkgs/development/python-modules/jaxopt/default.nix
@@ -25,6 +25,7 @@ buildPythonPackage (finalAttrs: {
   pname = "jaxopt";
   version = "0.8.5";
   pyproject = true;
+  __structuredAttrs = true;
 
   src = fetchFromGitHub {
     owner = "google";
@@ -88,6 +89,9 @@ buildPythonPackage (finalAttrs: {
     # https://github.com/google/jaxopt/issues/618
     "test_binary_logit_log_likelihood"
 
+    # AssertionError: Array(False, dtype=bool) is not true (since jax 0.12.0)
+    "test_hager_zhang_linesearch_non_finite1"
+
     # AssertionError (flaky numerical tests)
     "test_Rosenbrock2"
     "test_Rosenbrock5"
diff --git a/pkgs/development/python-modules/numpyro/default.nix b/pkgs/development/python-modules/numpyro/default.nix
index 1f485bbd3b87..ca0598896b3f 100644
--- a/pkgs/development/python-modules/numpyro/default.nix
+++ b/pkgs/development/python-modules/numpyro/default.nix
@@ -102,6 +102,7 @@ buildPythonPackage (finalAttrs: {
     "test_cpu"
     "test_entropy_categorical"
     "test_gaussian_model"
+    "test_zero_inflated_logits_probs_agree"
 
     # >       with pytest.warns(UserWarning, match="Hessian of log posterior"):
     # E       Failed: DID NOT WARN. No warnings of type (<class 'UserWarning'>,) were emitted.
diff --git a/pkgs/development/python-modules/waymax/default.nix b/pkgs/development/python-modules/waymax/default.nix
index 6854193d72d1..884a07efbb70 100644
--- a/pkgs/development/python-modules/waymax/default.nix
+++ b/pkgs/development/python-modules/waymax/default.nix
@@ -38,6 +38,18 @@ buildPythonPackage (finalAttrs: {
     hash = "sha256-vS0ZxLPKcjRzu5QaFApOVYNO5LBNQjop8wzmdEVdMC8=";
   };
 
+  # jax>=0.12 no longer supports 2D vectors in jnp.cross
+  # ValueError: Dimension must be 3 for cross product
+  postPatch = ''
+    substituteInPlace waymax/metrics/roadgraph.py \
+      --replace-fail \
+        "jnp.cross(points_to_edge, nearest_vector_xys)" \
+        "(points_to_edge[..., 0] * nearest_vector_xys[..., 1] - points_to_edge[..., 1] * nearest_vector_xys[..., 0])" \
+      --replace-fail \
+        "jnp.cross(points_to_edge, prior_vector_xys)" \
+        "(points_to_edge[..., 0] * prior_vector_xys[..., 1] - points_to_edge[..., 1] * prior_vector_xys[..., 0])"
+  '';
+
   build-system = [ setuptools ];
 
   dependencies = [