diff options
| -rw-r--r-- | pkgs/development/python-modules/bsuite/default.nix | 10 | ||||
| -rw-r--r-- | pkgs/development/python-modules/distrax/default.nix | 2 | ||||
| -rw-r--r-- | pkgs/development/python-modules/dm-haiku/default.nix | 101 | ||||
| -rw-r--r-- | pkgs/development/python-modules/dm-haiku/tests.nix | 68 | ||||
| -rw-r--r-- | pkgs/development/python-modules/rlax/default.nix | 2 |
5 files changed, 86 insertions, 97 deletions
diff --git a/pkgs/development/python-modules/bsuite/default.nix b/pkgs/development/python-modules/bsuite/default.nix index bf85d8fe7a090..3c11353bb1019 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,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 diff --git a/pkgs/development/python-modules/distrax/default.nix b/pkgs/development/python-modules/distrax/default.nix index 616dbae7a4fdd..cc667cc6bf19b 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 ]; diff --git a/pkgs/development/python-modules/dm-haiku/default.nix b/pkgs/development/python-modules/dm-haiku/default.nix index 08c1716867a7e..cb97e2f837afe 100644 --- a/pkgs/development/python-modules/dm-haiku/default.nix +++ b/pkgs/development/python-modules/dm-haiku/default.nix @@ -1,32 +1,39 @@ -{ buildPythonPackage +{ lib +, buildPythonPackage , fetchFromGitHub , fetchpatch -, callPackage -, lib +, absl-py +, flax +, jaxlib , jmp +, numpy , tabulate -, jaxlib +, 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 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 dec909729dcf1..0000000000000 --- a/pkgs/development/python-modules/dm-haiku/tests.nix +++ /dev/null @@ -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" - ]; - -} diff --git a/pkgs/development/python-modules/rlax/default.nix b/pkgs/development/python-modules/rlax/default.nix index 1bb7189d6957b..ceb8e9758619d 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 ]; |