pyrox.dev
1{ buildPythonPackage
2, fetchFromGitHub
3, fetchpatch
4, callPackage
5, lib
6, jmp
7, tabulate
8, jaxlib
9}:
10
11buildPythonPackage rec {
12 pname = "dm-haiku";
13 version = "0.0.10";
14
15 src = fetchFromGitHub {
16 owner = "deepmind";
17 repo = pname;
18 rev = "refs/tags/v${version}";
19 hash = "sha256-EZx3o6PgTeFjTwI9Ko9H39EqPSE0yLWWpsdqX6ALlo4=";
20 };
21
22 patches = [
23 # https://github.com/deepmind/dm-haiku/issues/717
24 (fetchpatch {
25 name = "remove-typing-extensions.patch";
26 url = "https://github.com/deepmind/dm-haiku/commit/c22867db1a3314a382bd2ce36511e2b756dc32a8.patch";
27 hash = "sha256-SxJc8FrImwMqTJ5OuJ1f4T+HfHgW/sGqXeIqlxEatlE=";
28 })
29 # https://github.com/deepmind/dm-haiku/pull/672
30 (fetchpatch {
31 name = "fix-find-namespace-packages.patch";
32 url = "https://github.com/deepmind/dm-haiku/commit/728031721f77d9aaa260bba0eddd9200d107ba5d.patch";
33 hash = "sha256-qV94TdJnphlnpbq+B0G3KTx5CFGPno+8FvHyu/aZeQE=";
34 })
35 ];
36
37 outputs = [
38 "out"
39 "testsout"
40 ];
41
42 propagatedBuildInputs = [
43 jaxlib
44 jmp
45 tabulate
46 ];
47
48 pythonImportsCheck = [
49 "haiku"
50 ];
51
52 postInstall = ''
53 mkdir $testsout
54 cp -R examples $testsout/examples
55 '';
56
57 # check in passthru.tests.pytest to escape infinite recursion with bsuite
58 doCheck = false;
59
60 passthru.tests = {
61 pytest = callPackage ./tests.nix { };
62 };
63
64 meta = with lib; {
65 description = "Haiku is a simple neural network library for JAX developed by some of the authors of Sonnet.";
66 homepage = "https://github.com/deepmind/dm-haiku";
67 license = licenses.asl20;
68 maintainers = with maintainers; [ ndl ];
69 };
70}