From 2cb66465f389a21c5d6fcc9876ac722846926b16 Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Tue, 25 Feb 2025 18:24:48 +0100
Subject: [PATCH] python312Packages.dm-haiku: cleanup & disable on python 3.13

---
 .../python-modules/dm-haiku/default.nix       | 22 ++++++++++++++-----
 1 file changed, 17 insertions(+), 5 deletions(-)

diff --git a/pkgs/development/python-modules/dm-haiku/default.nix b/pkgs/development/python-modules/dm-haiku/default.nix
index 96fbd446ecb4..5ea929f8570e 100644
--- a/pkgs/development/python-modules/dm-haiku/default.nix
+++ b/pkgs/development/python-modules/dm-haiku/default.nix
@@ -1,16 +1,25 @@
 {
   lib,
   buildPythonPackage,
+  pythonAtLeast,
   fetchFromGitHub,
   fetchpatch,
+
+  # build-system
   setuptools,
+
+  # dependencies
   absl-py,
-  flax,
-  jax,
   jaxlib,
   jmp,
   numpy,
   tabulate,
+
+  # optional-dependencies
+  jax,
+  flax,
+
+  # tests
   pytest-xdist,
   pytestCheckHook,
   bsuite,
@@ -30,6 +39,9 @@ let
     version = "0.0.13";
     pyproject = true;
 
+    # ImportError: `haiku.experimental.flax` features require `flax` to be installed.
+    disabled = pythonAtLeast "3.13";
+
     src = fetchFromGitHub {
       owner = "deepmind";
       repo = "dm-haiku";
@@ -117,11 +129,11 @@ let
       dontInstall = true;
     });
 
-    meta = with lib; {
+    meta = {
       description = "Haiku is a simple neural network library for JAX developed by some of the authors of Sonnet";
       homepage = "https://github.com/deepmind/dm-haiku";
-      license = licenses.asl20;
-      maintainers = with maintainers; [ ndl ];
+      license = lib.licenses.asl20;
+      maintainers = with lib.maintainers; [ ndl ];
     };
   };
 in