From 7d5f6f69f8b7e287f2644ef2da2f0987c883e850 Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Sat, 30 Aug 2025 11:15:32 +0200
Subject: [PATCH] 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
---
 .../python-modules/equinox/default.nix        | 22 ++++++++++++++++++-
 1 file changed, 21 insertions(+), 1 deletion(-)

diff --git a/pkgs/development/python-modules/equinox/default.nix b/pkgs/development/python-modules/equinox/default.nix
index 67b924287e3d..b56147a10b5b 100644
--- a/pkgs/development/python-modules/equinox/default.nix
+++ b/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 ];