From f79856aae4630f5e2063e37f6da86a201067decf Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Tue, 24 Mar 2026 08:01:18 +0000
Subject: [PATCH] python3Packages.clu: add missing chex test dependency

---
 pkgs/development/python-modules/clu/default.nix | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/pkgs/development/python-modules/clu/default.nix b/pkgs/development/python-modules/clu/default.nix
index 459c30bc0691..9f6f4f966517 100644
--- a/pkgs/development/python-modules/clu/default.nix
+++ b/pkgs/development/python-modules/clu/default.nix
@@ -19,6 +19,7 @@
   wrapt,
 
   # tests
+  chex,
   keras,
   pytestCheckHook,
   tensorflow,
@@ -26,7 +27,7 @@
   torch,
 }:
 
-buildPythonPackage rec {
+buildPythonPackage (finalAttrs: {
   pname = "clu";
   version = "0.0.12";
   pyproject = true;
@@ -34,7 +35,7 @@ buildPythonPackage rec {
   src = fetchFromGitHub {
     owner = "google";
     repo = "CommonLoopUtils";
-    tag = "v${version}";
+    tag = "v${finalAttrs.version}";
     hash = "sha256-ntqRz3fCXMf0EDQsddT68Mdi105ECBVQpVsStzk2kvQ=";
   };
 
@@ -59,6 +60,7 @@ buildPythonPackage rec {
   pythonImportsCheck = [ "clu" ];
 
   nativeCheckInputs = [
+    chex
     keras
     pytestCheckHook
     tensorflow
@@ -76,8 +78,8 @@ buildPythonPackage rec {
   meta = {
     description = "Common training loops in JAX";
     homepage = "https://github.com/google/CommonLoopUtils";
-    changelog = "https://github.com/google/CommonLoopUtils/blob/${src.tag}/CHANGELOG.md";
+    changelog = "https://github.com/google/CommonLoopUtils/blob/${finalAttrs.src.tag}/CHANGELOG.md";
     license = lib.licenses.asl20;
     maintainers = with lib.maintainers; [ GaetanLepage ];
   };
-}
+})