From 25f652d97518adf3fe4c265c5b4a5639736288cd Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Sun, 19 Apr 2026 22:00:23 +0000
Subject: [PATCH] python3Packages.chex: fix jax 0.10.0 compatibility

---
 pkgs/development/python-modules/chex/default.nix | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/pkgs/development/python-modules/chex/default.nix b/pkgs/development/python-modules/chex/default.nix
index 739fa4bcdc5b..ed93783889d6 100644
--- a/pkgs/development/python-modules/chex/default.nix
+++ b/pkgs/development/python-modules/chex/default.nix
@@ -2,6 +2,7 @@
   lib,
   buildPythonPackage,
   fetchFromGitHub,
+  fetchpatch,
 
   # build-system
   flit-core,
@@ -33,6 +34,18 @@ buildPythonPackage (finalAttrs: {
     hash = "sha256-lJ9+kvG7dRtfDVgvkcJ9/jtnX0lMfxY4mmZ290y/74U=";
   };
 
+  patches = [
+    # jax.device_put_replicated is removed in jax 0.10.0
+    # This fix was merged upstream -> remove when updating to the next release
+    (fetchpatch {
+      url = "https://github.com/google-deepmind/chex/commit/5fbd2c9a9936799daf92354e0307b9e88b9cc163.patch";
+      excludes = [
+        "chex/_src/variants.py"
+      ];
+      hash = "sha256-ZTimSq7/yt2UEiWmLcfFBadX8+VcaxuPhkQJEyiEZlE=";
+    })
+  ];
+
   build-system = [
     flit-core
   ];