@@ -0,0 +1,71 @@
1+{ buildPythonPackage
2+, cloudpickle
3+, dm-haiku
4+, einops
5+, fetchFromGitHub
6+, flax
7+, hypothesis
8+, keras
9+, lib
10+, poetry-core
11+, pytestCheckHook
12+, pyyaml
13+, rich
14+, tensorflow
15+, treeo
16+}:
17+
18+buildPythonPackage rec {
19+ pname = "treex";
20+ version = "0.6.7";
21+ format = "pyproject";
22+
23+ src = fetchFromGitHub {
24+ owner = "cgarciae";
25+ repo = pname;
26+ rev = version;
27+ sha256 = "1hl3wj71c7cp7jzkhyjy7xgs2vc8c89icq0bgfr49y4pwv69n43m";
28+ };
29+
30+ patches = [
31+ ./relax-deps.patch
32+ ];
33+
34+ nativeBuildInputs = [
35+ poetry-core
36+ ];
37+
38+ propagatedBuildInputs = [
39+ einops
40+ flax
41+ pyyaml
42+ rich
43+ treeo
44+ ];
45+
46+ checkInputs = [
47+ cloudpickle
48+ dm-haiku
49+ hypothesis
50+ keras
51+ pytestCheckHook
52+ tensorflow
53+ ];
54+
55+ pythonImportsCheck = [
56+ "treex"
57+ ];
58+
59+ disabledTestPaths = [
60+ # Require `torchmetrics` which is not packaged in `nixpkgs`.
61+ "tests/metrics/test_mean_absolute_error.py"
62+ "tests/metrics/test_mean_square_error.py"
63+ ];
64+
65+ meta = with lib; {
66+ description = "Pytree Module system for Deep Learning in JAX";
67+ homepage = "https://github.com/cgarciae/treex";
68+ license = licenses.mit;
69+ maintainers = with maintainers; [ ndl ];
70+ };
71+}