From 0e1746ad578321d7a977dce44072341dd9bf04fd Mon Sep 17 00:00:00 2001 From: Gaetan Lepage Date: Tue, 9 Jan 2024 21:54:32 +0100 Subject: [PATCH 1/6] python311Packages.dm-haiku: 0.10.0 -> 0.11.0 Changelog: https://github.com/google-deepmind/dm-haiku/releases/tag/v0.0.11 --- .../python-modules/dm-haiku/default.nix | 18 +++++++++--------- .../python-modules/dm-haiku/tests.nix | 1 + 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/pkgs/development/python-modules/dm-haiku/default.nix b/pkgs/development/python-modules/dm-haiku/default.nix index 08c1716867a7..6ff6ff412b7e 100644 --- a/pkgs/development/python-modules/dm-haiku/default.nix +++ b/pkgs/development/python-modules/dm-haiku/default.nix @@ -1,6 +1,9 @@ { buildPythonPackage , fetchFromGitHub , fetchpatch +, absl-py +, flax +, numpy , callPackage , lib , jmp @@ -10,23 +13,17 @@ 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"; @@ -41,8 +38,11 @@ buildPythonPackage rec { ]; propagatedBuildInputs = [ + absl-py + flax jaxlib jmp + numpy tabulate ]; diff --git a/pkgs/development/python-modules/dm-haiku/tests.nix b/pkgs/development/python-modules/dm-haiku/tests.nix index dec909729dcf..3a99bd3ac85a 100644 --- a/pkgs/development/python-modules/dm-haiku/tests.nix +++ b/pkgs/development/python-modules/dm-haiku/tests.nix @@ -20,6 +20,7 @@ buildPythonPackage { pname = "dm-haiku-tests"; + format = "other"; inherit (dm-haiku) version; src = dm-haiku.testsout; From ee1c02123689dbcf322e1074a0d9aab55b50d1e4 Mon Sep 17 00:00:00 2001 From: Gaetan Lepage Date: Tue, 9 Jan 2024 22:04:54 +0100 Subject: [PATCH 2/6] python311Packages.bsuite: use propagatedBuildInputs for dependencies --- pkgs/development/python-modules/bsuite/default.nix | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pkgs/development/python-modules/bsuite/default.nix b/pkgs/development/python-modules/bsuite/default.nix index bf85d8fe7a09..f6560941f3d6 100644 --- a/pkgs/development/python-modules/bsuite/default.nix +++ b/pkgs/development/python-modules/bsuite/default.nix @@ -35,7 +35,7 @@ let bsuite = buildPythonPackage rec { hash = "sha256-ak9McvXl7Nz5toUaPaRaJek9lurxiQiIW209GnZEjX0="; }; - buildInputs = [ + propagatedBuildInputs = [ absl-py dm-env dm-tree From 53b8616b2d1d608c1c07ea0c883112dfc37a8a01 Mon Sep 17 00:00:00 2001 From: Gaetan Lepage Date: Wed, 10 Jan 2024 09:16:38 +0100 Subject: [PATCH 3/6] python311Packages.bsuite: include patch that fixes deprecated use of np.int --- pkgs/development/python-modules/bsuite/default.nix | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/pkgs/development/python-modules/bsuite/default.nix b/pkgs/development/python-modules/bsuite/default.nix index f6560941f3d6..3c11353bb101 100644 --- a/pkgs/development/python-modules/bsuite/default.nix +++ b/pkgs/development/python-modules/bsuite/default.nix @@ -1,6 +1,7 @@ { lib , fetchPypi , buildPythonPackage +, fetchpatch , frozendict , termcolor , matplotlib @@ -35,6 +36,13 @@ let bsuite = buildPythonPackage rec { hash = "sha256-ak9McvXl7Nz5toUaPaRaJek9lurxiQiIW209GnZEjX0="; }; + 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 From 175f71a0f352df6d5ef98816ff22245704674787 Mon Sep 17 00:00:00 2001 From: Gaetan Lepage Date: Wed, 10 Jan 2024 09:30:56 +0100 Subject: [PATCH 4/6] python311Packages.dm-haiku: fix tests --- .../python-modules/dm-haiku/default.nix | 87 ++++++++++++++----- .../python-modules/dm-haiku/tests.nix | 69 --------------- 2 files changed, 66 insertions(+), 90 deletions(-) delete mode 100644 pkgs/development/python-modules/dm-haiku/tests.nix diff --git a/pkgs/development/python-modules/dm-haiku/default.nix b/pkgs/development/python-modules/dm-haiku/default.nix index 6ff6ff412b7e..cb97e2f837af 100644 --- a/pkgs/development/python-modules/dm-haiku/default.nix +++ b/pkgs/development/python-modules/dm-haiku/default.nix @@ -1,17 +1,27 @@ -{ buildPythonPackage +{ lib +, buildPythonPackage , fetchFromGitHub , fetchpatch , absl-py , flax -, numpy -, callPackage -, lib -, jmp -, tabulate , 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.11"; format = "setuptools"; @@ -32,11 +42,6 @@ buildPythonPackage rec { }) ]; - outputs = [ - "out" - "testsout" - ]; - propagatedBuildInputs = [ absl-py flax @@ -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 diff --git a/pkgs/development/python-modules/dm-haiku/tests.nix b/pkgs/development/python-modules/dm-haiku/tests.nix deleted file mode 100644 index 3a99bd3ac85a..000000000000 --- a/pkgs/development/python-modules/dm-haiku/tests.nix +++ /dev/null @@ -1,69 +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"; - format = "other"; - 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" - ]; - -} From 6fe5e5a54ac122f9f7f3235fac8e82e28a286638 Mon Sep 17 00:00:00 2001 From: Gaetan Lepage Date: Wed, 10 Jan 2024 11:40:12 +0100 Subject: [PATCH 5/6] python311Packages.distrax: use pytest-xdist --- pkgs/development/python-modules/distrax/default.nix | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pkgs/development/python-modules/distrax/default.nix b/pkgs/development/python-modules/distrax/default.nix index 616dbae7a4fd..cc667cc6bf19 100644 --- a/pkgs/development/python-modules/distrax/default.nix +++ b/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 ]; From c92778b9e5a5d23a1f5bab3b987fc9577eaff43d Mon Sep 17 00:00:00 2001 From: Gaetan Lepage Date: Wed, 10 Jan 2024 11:50:07 +0100 Subject: [PATCH 6/6] python311Packages.rlax: use pytest-xdist --- pkgs/development/python-modules/rlax/default.nix | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pkgs/development/python-modules/rlax/default.nix b/pkgs/development/python-modules/rlax/default.nix index 1bb7189d6957..ceb8e9758619 100644 --- a/pkgs/development/python-modules/rlax/default.nix +++ b/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 ];