From 189daa5eb3288dd77e7d72805ed677d776705360 Mon Sep 17 00:00:00 2001
From: Ben Darwin <bcdarwin@gmail.com>
Date: Sun, 13 Aug 2023 12:10:07 -0400
Subject: [PATCH] python310Packages.chex: 0.1.6 -> 0.1.82; unbreak

---
 .../python-modules/chex/default.nix           | 32 ++++++++-----------
 1 file changed, 13 insertions(+), 19 deletions(-)

diff --git a/pkgs/development/python-modules/chex/default.nix b/pkgs/development/python-modules/chex/default.nix
index b93d3baddcc4..7d7912ff0f45 100644
--- a/pkgs/development/python-modules/chex/default.nix
+++ b/pkgs/development/python-modules/chex/default.nix
@@ -1,35 +1,39 @@
-{ absl-py
+{ lib
 , buildPythonPackage
+, pythonOlder
+, fetchFromGitHub
+, absl-py
 , cloudpickle
 , dm-tree
-, fetchFromGitHub
 , jax
 , jaxlib
-, lib
 , numpy
 , pytestCheckHook
 , toolz
+, typing-extensions
 }:
 
 buildPythonPackage rec {
   pname = "chex";
-  version = "0.1.6";
+  version = "0.1.82";
   format = "setuptools";
 
+  disabled = pythonOlder "3.9";
+
   src = fetchFromGitHub {
     owner = "deepmind";
     repo = pname;
     rev = "refs/tags/v${version}";
-    hash = "sha256-VolRlLLgKga9S17ByVrYya9VPtu9yiOnvt/WmlE1DOc=";
+    hash = "sha256-xBq22AaR2Tp1NSPefEyvCDeUYqRZlAf5LVHWo0luiXk=";
   };
 
   propagatedBuildInputs = [
     absl-py
-    cloudpickle
-    dm-tree
+    jaxlib
     jax
     numpy
     toolz
+    typing-extensions
   ];
 
   pythonImportsCheck = [
@@ -37,21 +41,11 @@ buildPythonPackage rec {
   ];
 
   nativeCheckInputs = [
-    jaxlib
+    cloudpickle
+    dm-tree
     pytestCheckHook
   ];
 
-  disabledTests = [
-    # See https://github.com/deepmind/chex/issues/204.
-    "test_uninspected_checks"
-
-    # These tests started failing at some point after upgrading to 0.1.5
-    "test_useful_failure"
-    "TreeAssertionsTest"
-    "PmapFakeTest"
-    "WithDeviceTest"
-  ];
-
   meta = with lib; {
     description = "Chex is a library of utilities for helping to write reliable JAX code.";
     homepage = "https://github.com/deepmind/chex";