From 16776b808c64a3efbf5fc2fa1a7331170a23252f Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Fri, 12 Jun 2026 15:18:44 +0000
Subject: [PATCH] python3Packages.chex: 0.1.91 -> 0.1.92

Diff: https://github.com/google-deepmind/chex/compare/v0.1.91...v0.1.92

Changelog: https://github.com/google-deepmind/chex/releases/tag/v0.1.92
---
 .../python-modules/chex/default.nix           | 31 +++----------------
 1 file changed, 4 insertions(+), 27 deletions(-)

diff --git a/pkgs/development/python-modules/chex/default.nix b/pkgs/development/python-modules/chex/default.nix
index ed93783889d6..bdd741606efe 100644
--- a/pkgs/development/python-modules/chex/default.nix
+++ b/pkgs/development/python-modules/chex/default.nix
@@ -2,7 +2,6 @@
   lib,
   buildPythonPackage,
   fetchFromGitHub,
-  fetchpatch,
 
   # build-system
   flit-core,
@@ -23,29 +22,17 @@
 
 buildPythonPackage (finalAttrs: {
   pname = "chex";
-  version = "0.1.91";
+  version = "0.1.92";
   pyproject = true;
   __structuredAttrs = true;
 
   src = fetchFromGitHub {
-    owner = "deepmind";
+    owner = "google-deepmind";
     repo = "chex";
     tag = "v${finalAttrs.version}";
-    hash = "sha256-lJ9+kvG7dRtfDVgvkcJ9/jtnX0lMfxY4mmZ290y/74U=";
+    hash = "sha256-PM76Q72Bgyms7dROJkmlpPuDvtqjHLPTDkUYqo08T74=";
   };
 
-  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
   ];
@@ -67,19 +54,9 @@ buildPythonPackage (finalAttrs: {
     pytestCheckHook
   ];
 
-  disabledTests = [
-    # Jax 0.8.2 incompatibility (reported at https://github.com/google-deepmind/chex/issues/422)
-    # AssertionError: AssertionError not raised
-    "test_assert_tree_is_on_device"
-    # AssertionError: "\[Chex\]\ [\s\S]*sharded arrays are disallowed" does not match ...
-    "test_assert_tree_is_on_host"
-    # AssertionError: [Chex] Assertion assert_tree_is_sharded failed: ...
-    "test_assert_tree_is_sharded"
-  ];
-
   meta = {
     description = "Library of utilities for helping to write reliable JAX code";
-    homepage = "https://github.com/deepmind/chex";
+    homepage = "https://github.com/google-deepmind/chex";
     changelog = "https://github.com/google-deepmind/chex/releases/tag/${finalAttrs.src.tag}";
     license = lib.licenses.asl20;
     maintainers = with lib.maintainers; [ ndl ];