From 368f097532983ecc9340ae170ce815eddd70ff72 Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Mon, 20 Oct 2025 07:40:58 +0000
Subject: [PATCH] python3Packages.dm-haiku: 0.0.13 -> 0.0.15

Diff: https://github.com/deepmind/dm-haiku/compare/v0.0.13...v0.0.15

Changelog:
https://github.com/google-deepmind/dm-haiku/releases/tag/v0.0.15
---
 .../python-modules/dm-haiku/default.nix       | 39 ++++++++++---------
 1 file changed, 20 insertions(+), 19 deletions(-)

diff --git a/pkgs/development/python-modules/dm-haiku/default.nix b/pkgs/development/python-modules/dm-haiku/default.nix
index 6e39b1e872ff..51745dfdc89a 100644
--- a/pkgs/development/python-modules/dm-haiku/default.nix
+++ b/pkgs/development/python-modules/dm-haiku/default.nix
@@ -36,17 +36,14 @@
 let
   dm-haiku = buildPythonPackage rec {
     pname = "dm-haiku";
-    version = "0.0.13";
+    version = "0.0.15";
     pyproject = true;
 
-    # ImportError: `haiku.experimental.flax` features require `flax` to be installed.
-    disabled = pythonAtLeast "3.13";
-
     src = fetchFromGitHub {
       owner = "deepmind";
       repo = "dm-haiku";
       tag = "v${version}";
-      hash = "sha256-RJpQ9BzlbQ4X31XoJFnsZASiaC9fP2AdyuTAGINhMxs=";
+      hash = "sha256-phJ0f+effHQzuAVtPBR0bY3C0p//LBY7k1ci4mXBGfU=";
     };
 
     patches = [
@@ -58,17 +55,6 @@ let
       })
     ];
 
-    # AttributeError: jax.core.Var was removed in JAX v0.6.0. Use jax.extend.core.Var instead, and
-    # see https://docs.jax.dev/en/latest/jax.extend.html for details.
-    # Already on master: https://github.com/google-deepmind/dm-haiku/commit/cfe8480d253a93100bf5e2d24c40435a95399c96
-    # TODO: remove at the next release
-    postPatch = ''
-      substituteInPlace haiku/_src/jaxpr_info.py \
-        --replace-fail "jax.core.JaxprEqn" "jax.extend.core.JaxprEqn" \
-        --replace-fail "jax.core.Var" "jax.extend.core.Var" \
-        --replace-fail "jax.core.Jaxpr" "jax.extend.core.Jaxpr"
-    '';
-
     build-system = [ setuptools ];
 
     dependencies = [
@@ -97,11 +83,12 @@ let
       dm-env
       dm-haiku
       dm-tree
+      flax
       jaxlib
       optax
       pytest-xdist
       pytestCheckHook
-      rlax
+      # rlax (broken dependency tensorflow-probability)
       tensorflow
     ];
 
@@ -120,11 +107,24 @@ let
       # other tests to fail.
       # https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#double-64bit-precision
       "test_doctest_haiku.experimental"
+
+      # AssertionError: 1 != 0 : 1 doctests failed
+      "test_doctest_haiku"
+
+      # ValueError: pmap wrapped function must be passed at least one argument containing an array,
+      # got empty *args=() and **kwargs={}
+      "test_equivalent_when_passing_transformed_fn2"
+
+      # AssertionError: ValueError not raised
+      "test_passing_function_to_transform_pmap_transform"
+      "test_passing_function_to_transform_pmap_transform_with_state"
     ];
 
     disabledTestPaths = [
-      # Those tests requires a more recent version of tensorflow. The current one (2.13) is not enough.
-      "haiku/_src/integration/jax2tf_test.py"
+      # Require rlax which is unavailable as its dependency tensorflow-probability is broken
+      "examples/impala/actor_test.py"
+      "examples/impala/learner_test.py"
+      "examples/impala_lite_test.py"
     ];
 
     doCheck = false;
@@ -143,6 +143,7 @@ let
     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";
+      changelog = "https://github.com/google-deepmind/dm-haiku/releases/tag/${src.tag}";
       license = lib.licenses.asl20;
       maintainers = with lib.maintainers; [ ndl ];
     };