python3Packages.equinox: fix jax 0.8.2 compatibility

https://github.com/patrick-kidger/equinox/pull/1162
2025-12-28 17:03:58 +00:00
parent c4a987ff82
commit 70e989a146
1 changed files with 6 additions and 0 deletions
@@ -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 ];