python3Packages.equinox: backport patches to fix with jax 0.7.0

https://github.com/patrick-kidger/equinox/pull/1086
https://github.com/patrick-kidger/equinox/pull/1082

pkgs/development/python-modules/equinox/default.nix

@@ -3,6 +3,7 @@
   stdenv,
   buildPythonPackage,
   fetchFromGitHub,
+  fetchpatch2,
 
   # build-system
   hatchling,
@@ -32,12 +33,31 @@ buildPythonPackage rec {
     hash = "sha256-zXgAuFGWKHShKodi9swnWIry4VU9s4pBhBRoK5KzaL0=";
   };
 
+  patches = [
+    # The following two patches have been merged upstream and should be removed when updating to the next release
+    # They fix the incompatibilities with jax>=0.7.0
+
+    # https://github.com/patrick-kidger/equinox/pull/1086
+    (fetchpatch2 {
+      name = "remove-deprecated-batching-NotMapped";
+      url = "https://github.com/patrick-kidger/equinox/commit/6a6a441ced2fe64191a087752f1c2e71a6ce39f1.patch";
+      hash = "sha256-tzHFjMI3gAIh5MPkdbmzsky/oFjDEbOIkPGQMQ+gcQQ=";
+    })
+
+    # https://github.com/patrick-kidger/equinox/pull/1082
+    (fetchpatch2 {
+      name = "allow-creating-weak-references-to-flatten";
+      url = "https://github.com/patrick-kidger/equinox/commit/62b3c94ad56bdb63524702b320e977d2d93dbe72.patch";
+      hash = "sha256-c1FKCnC3/okuP2VJV4h7sPRYQeYJZSdzEG5ETL2M35k=";
+    })
+  ];
+
   # Relax speed constraints on tests that can fail on busy builders
   postPatch = ''
     substituteInPlace tests/test_while_loop.py \
       --replace-fail "speed < 0.1" "speed < 0.5" \
       --replace-fail "speed < 0.5" "speed < 1" \
-      --replace-fail "speed < 1" "speed < 4" \
+      --replace-fail "speed < 1" "speed < 20"
   '';
 
   build-system = [ hatchling ];