From 70e989a146b5f8dccb1c95171e08a314dee9b05e Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Fri, 26 Dec 2025 23:48:26 +0000
Subject: [PATCH] python3Packages.equinox: fix jax 0.8.2 compatibility

https://github.com/patrick-kidger/equinox/pull/1162
---
 pkgs/development/python-modules/equinox/default.nix | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/pkgs/development/python-modules/equinox/default.nix b/pkgs/development/python-modules/equinox/default.nix
index 0ede965dc4e7..29337297d713 100644
--- a/pkgs/development/python-modules/equinox/default.nix
+++ b/pkgs/development/python-modules/equinox/default.nix
@@ -39,6 +39,12 @@ buildPythonPackage rec {
       --replace-fail "speed < 0.5" "speed < 1" \
       --replace-fail "speed < 1" "speed < 20" \
       --replace-fail "speed < 2" "speed < 20"
+  ''
+  # Fix jax 0.8.2 compat
+  # Fix submitted upstream: https://github.com/patrick-kidger/equinox/pull/1162
+  + ''
+    substituteInPlace equinox/_ad.py equinox/internal/_primitive.py \
+      --replace-fail "jax.core.get_aval(" "jax.typeof("
   '';
 
   build-system = [ hatchling ];