Merge pull request #279901 from GaetanLepage/dm-haiku

python311Packages.dm-haiku: 0.10.0 -> 0.11.0

pkgs/development/python-modules/bsuite/default.nix

@@ -1,6 +1,7 @@
 { lib
 , fetchPypi
 , buildPythonPackage
+, fetchpatch
 , frozendict
 , termcolor
 , matplotlib
@@ -35,7 +36,14 @@ let bsuite = buildPythonPackage rec {
     hash = "sha256-ak9McvXl7Nz5toUaPaRaJek9lurxiQiIW209GnZEjX0=";
   };
 
-  buildInputs = [
+  patches = [
+    (fetchpatch {  # Convert np.int -> np.int32 since np.int is deprecated (https://github.com/google-deepmind/bsuite/pull/48)
+      url = "https://github.com/google-deepmind/bsuite/pull/48/commits/f8d81b2f1c27ef2c8c71ae286001ed879ea306ab.patch";
+      hash = "sha256-FXtvVS+U8brulq8Z27+yWIimB+kigGiUOIv1SHb1TA8=";
+    })
+  ];
+
+  propagatedBuildInputs = [
     absl-py
     dm-env
     dm-tree

pkgs/development/python-modules/distrax/default.nix

@@ -7,6 +7,7 @@
 , numpy
 , tensorflow-probability
 , dm-haiku
+, pytest-xdist
 , pytestCheckHook
 }:
 
@@ -33,6 +34,7 @@ buildPythonPackage rec {
 
   nativeCheckInputs = [
     dm-haiku
+    pytest-xdist
     pytestCheckHook
   ];
 

pkgs/development/python-modules/dm-haiku/default.nix

@@ -1,32 +1,39 @@
-{ buildPythonPackage
+{ lib
+, buildPythonPackage
 , fetchFromGitHub
 , fetchpatch
-, callPackage
-, lib
-, jmp
-, tabulate
+, absl-py
+, flax
 , jaxlib
+, jmp
+, numpy
+, tabulate
+, pytest-xdist
+, pytestCheckHook
+, bsuite
+, chex
+, cloudpickle
+, dill
+, dm-env
+, dm-tree
+, optax
+, rlax
+, tensorflow
 }:
 
-buildPythonPackage rec {
+let dm-haiku = buildPythonPackage rec {
   pname = "dm-haiku";
-  version = "0.0.10";
+  version = "0.0.11";
   format = "setuptools";
 
   src = fetchFromGitHub {
     owner = "deepmind";
-    repo = pname;
+    repo = "dm-haiku";
     rev = "refs/tags/v${version}";
-    hash = "sha256-EZx3o6PgTeFjTwI9Ko9H39EqPSE0yLWWpsdqX6ALlo4=";
+    hash = "sha256-xve1vNsVOC6/HVtzmzswM/Sk3uUNaTtqNAKheFb/tmI=";
   };
 
   patches = [
-    # https://github.com/deepmind/dm-haiku/issues/717
-    (fetchpatch {
-      name = "remove-typing-extensions.patch";
-      url = "https://github.com/deepmind/dm-haiku/commit/c22867db1a3314a382bd2ce36511e2b756dc32a8.patch";
-      hash = "sha256-SxJc8FrImwMqTJ5OuJ1f4T+HfHgW/sGqXeIqlxEatlE=";
-    })
     # https://github.com/deepmind/dm-haiku/pull/672
     (fetchpatch {
       name = "fix-find-namespace-packages.patch";
@@ -35,14 +42,12 @@ buildPythonPackage rec {
     })
   ];
 
-  outputs = [
-    "out"
-    "testsout"
-  ];
-
   propagatedBuildInputs = [
+    absl-py
+    flax
     jaxlib
     jmp
+    numpy
     tabulate
   ];
 
@@ -50,17 +55,56 @@ buildPythonPackage rec {
     "haiku"
   ];
 
-  postInstall = ''
-    mkdir $testsout
-    cp -R examples $testsout/examples
-  '';
+  nativeCheckInputs = [
+    bsuite
+    chex
+    cloudpickle
+    dill
+    dm-env
+    dm-haiku
+    dm-tree
+    jaxlib
+    optax
+    pytest-xdist
+    pytestCheckHook
+    rlax
+    tensorflow
+  ];
+
+  disabledTests = [
+    # See https://github.com/deepmind/dm-haiku/issues/366.
+    "test_jit_Recurrent"
+
+    # Assertion errors
+    "testShapeChecking0"
+    "testShapeChecking1"
+
+    # This test requires a more recent version of tensorflow. The current one (2.13) is not enough.
+    "test_reshape_convert"
+
+    # This test requires JAX support for double precision (64bit), but enabling this causes several
+    # other tests to fail.
+    # https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#double-64bit-precision
+    "test_doctest_haiku.experimental"
+  ];
+
+  disabledTestPaths = [
+    # Those tests requires a more recent version of tensorflow. The current one (2.13) is not enough.
+    "haiku/_src/integration/jax2tf_test.py"
+  ];
 
-  # check in passthru.tests.pytest to escape infinite recursion with bsuite
   doCheck = false;
 
-  passthru.tests = {
-    pytest = callPackage ./tests.nix { };
-  };
+  # check in passthru.tests.pytest to escape infinite recursion with bsuite
+  passthru.tests.pytest = dm-haiku.overridePythonAttrs (_: {
+    pname = "${pname}-tests";
+    doCheck = true;
+
+    # We don't have to install because the only purpose
+    # of this passthru test is to, well, test.
+    # This fixes having to set `catchConflicts` to false.
+    dontInstall = true;
+  });
 
   meta = with lib; {
     description = "Haiku is a simple neural network library for JAX developed by some of the authors of Sonnet.";
@@ -68,4 +112,5 @@ buildPythonPackage rec {
     license = licenses.asl20;
     maintainers = with maintainers; [ ndl ];
   };
-}
+};
+in dm-haiku

pkgs/development/python-modules/dm-haiku/tests.nix

@@ -1,68 +0,0 @@
-{ buildPythonPackage
-, dm-haiku
-, chex
-, cloudpickle
-, dill
-, dm-tree
-, jaxlib
-, pytest-xdist
-, pytestCheckHook
-, tensorflow
-, bsuite
-, frozendict
-, dm-env
-, scikit-image
-, rlax
-, distrax
-, tensorflow-probability
-, optax
-}:
-
-buildPythonPackage {
-  pname = "dm-haiku-tests";
-  inherit (dm-haiku) version;
-
-  src = dm-haiku.testsout;
-
-  dontBuild = true;
-  dontInstall = true;
-
-  nativeCheckInputs = [
-    bsuite
-    chex
-    cloudpickle
-    dill
-    distrax
-    dm-env
-    dm-haiku
-    dm-tree
-    frozendict
-    jaxlib
-    pytest-xdist
-    pytestCheckHook
-    optax
-    rlax
-    scikit-image
-    tensorflow
-    tensorflow-probability
-  ];
-
-  disabledTests = [
-    # See https://github.com/deepmind/dm-haiku/issues/366.
-    "test_jit_Recurrent"
-    # Assertion errors
-    "test_connect_conv_padding_function_same0"
-    "test_connect_conv_padding_function_valid0"
-    "test_connect_conv_padding_function_same1"
-    "test_connect_conv_padding_function_same2"
-    "test_connect_conv_padding_function_valid1"
-    "test_connect_conv_padding_function_valid2"
-    "test_invalid_axis_ListString"
-    "test_invalid_axis_String"
-    "test_simple_case"
-    "test_simple_case_with_scale"
-    "test_slice_axis"
-    "test_zero_inputs"
-  ];
-
-}

pkgs/development/python-modules/rlax/default.nix

@@ -12,6 +12,7 @@
 , tensorflow-probability
 , dm-haiku
 , optax
+, pytest-xdist
 , pytestCheckHook
 }:
 
@@ -49,6 +50,7 @@ buildPythonPackage rec {
   nativeCheckInputs = [
     dm-haiku
     optax
+    pytest-xdist
     pytestCheckHook
   ];