diff --git a/pkgs/development/python-modules/jax/default.nix b/pkgs/development/python-modules/jax/default.nix
index dfc221a76e56..4cd12283ab1b 100644
--- a/pkgs/development/python-modules/jax/default.nix
+++ b/pkgs/development/python-modules/jax/default.nix
@@ -6,6 +6,7 @@
   lapack,
   buildPythonPackage,
   fetchFromGitHub,
+  fetchpatch2,
   cudaSupport ? config.cudaSupport,
 
   # build-system
@@ -51,6 +52,14 @@ buildPythonPackage rec {
     hash = "sha256-Ilcp4WW65SyqrqDGBRdnB25m7OCbrsfdtxWvl0uTjNw=";
   };
 
+  patches = [
+    # https://github.com/jax-ml/jax/pull/32840
+    (fetchpatch2 {
+      url = "https://github.com/Prince213/jax/commit/af5c211d49f3b99447db2252d2cc2b8e0fb54d1c.patch?full_index=1";
+      hash = "sha256-ijEd+MDe91qyYfE+aMzR5rNmTeGadin6Io8PIfJWc3o=";
+    })
+  ];
+
   build-system = [ setuptools ];
 
   # The version is automatically set to ".dev" if this variable is not set.