1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
|
{
lib,
buildPythonPackage,
fetchFromGitHub,
fetchpatch,
absl-py,
flax,
jaxlib,
jmp,
numpy,
tabulate,
pytest-xdist,
pytestCheckHook,
bsuite,
chex,
cloudpickle,
dill,
dm-env,
dm-tree,
optax,
rlax,
tensorflow,
}:
let
dm-haiku = buildPythonPackage rec {
pname = "dm-haiku";
version = "0.0.12";
format = "setuptools";
src = fetchFromGitHub {
owner = "deepmind";
repo = "dm-haiku";
rev = "refs/tags/v${version}";
hash = "sha256-aJRXlMq4CNMH3ZSTDP8MgnVltdSc8l5raw4//KccL48=";
};
patches = [
# https://github.com/deepmind/dm-haiku/pull/672
(fetchpatch {
name = "fix-find-namespace-packages.patch";
url = "https://github.com/deepmind/dm-haiku/commit/728031721f77d9aaa260bba0eddd9200d107ba5d.patch";
hash = "sha256-qV94TdJnphlnpbq+B0G3KTx5CFGPno+8FvHyu/aZeQE=";
})
];
propagatedBuildInputs = [
absl-py
flax
jaxlib
jmp
numpy
tabulate
];
pythonImportsCheck = [ "haiku" ];
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"
];
doCheck = false;
# 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";
homepage = "https://github.com/deepmind/dm-haiku";
license = licenses.asl20;
maintainers = with maintainers; [ ndl ];
};
};
in
dm-haiku
|
