diff --git a/pkgs/development/python-modules/dm-haiku/default.nix b/pkgs/development/python-modules/dm-haiku/default.nix new file mode 100644 index 0000000000000..5468776c72d11 --- /dev/null +++ b/pkgs/development/python-modules/dm-haiku/default.nix @@ -0,0 +1,59 @@ +{ buildPythonPackage +, chex +, cloudpickle +, dill +, dm-tree +, fetchFromGitHub +, jmp +, lib +, pytestCheckHook +, tabulate +, tensorflow +}: + +buildPythonPackage rec { + pname = "dm-haiku"; + version = "0.0.5"; + + src = fetchFromGitHub { + owner = "deepmind"; + repo = pname; + rev = "v${version}"; + sha256 = "1mdqjcka0m1div63ngba8w8z94id4c1h8xqmnq1xpmgkc79224wa"; + }; + + propagatedBuildInputs = [ + jmp + tabulate + ]; + + checkInputs = [ + chex + cloudpickle + dm-tree + pytestCheckHook + tensorflow + ]; + + pythonImportsCheck = [ + "haiku" + ]; + + disabledTestPaths = [ + # These tests require `bsuite` which isn't packaged in `nixpkgs`. + "examples/impala_lite_test.py" + "examples/impala/actor_test.py" + "examples/impala/learner_test.py" + # This test breaks on multiple cases with TF-related errors, + # likely that's the reason the upstream uses TF-nightly for tests? + # `nixpkgs` doesn't have the corresponding TF version packaged. + "haiku/_src/integration/jax2tf_test.py" + ]; + + meta = with lib; { + 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 ]; + }; +} diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index fafac2e8f2462..78b5ee110f69b 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -2325,6 +2325,8 @@ in { dmenu-python = callPackage ../development/python-modules/dmenu { }; + dm-haiku = callPackage ../development/python-modules/dm-haiku { }; + dm-sonnet = callPackage ../development/python-modules/dm-sonnet { }; dm-tree = callPackage ../development/python-modules/dm-tree { };