pyrox.dev
1{ buildPythonPackage
2, cloudpickle
3, dm-haiku
4, einops
5, fetchFromGitHub
6, flax
7, hypothesis
8, jaxlib
9, keras
10, lib
11, poetry-core
12, pytestCheckHook
13, pyyaml
14, rich
15, tensorflow
16, treeo
17, torchmetrics
18, pythonRelaxDepsHook
19, torch
20}:
21
22buildPythonPackage rec {
23 pname = "treex";
24 version = "0.6.10";
25 format = "pyproject";
26
27 src = fetchFromGitHub {
28 owner = "cgarciae";
29 repo = pname;
30 rev = version;
31 hash = "sha256-ZHfgmRNbFh8DFZkmilY0pmRNQhJFqT689I7Lu8FuFm4=";
32 };
33
34 # At the time of writing (2022-03-29), rich is currently at version 11.0.0.
35 # The treeo dependency is compatible with a patch, but not marked as such in
36 # treex. See https://github.com/cgarciae/treex/issues/68.
37 pythonRelaxDeps = [
38 "certifi"
39 "flax"
40 "rich"
41 "treeo"
42 ];
43
44 nativeBuildInputs = [
45 poetry-core
46 pythonRelaxDepsHook
47 ];
48
49 buildInputs = [ jaxlib ];
50
51 propagatedBuildInputs = [
52 einops
53 flax
54 pyyaml
55 rich
56 treeo
57 torch
58 ];
59
60 checkInputs = [
61 cloudpickle
62 dm-haiku
63 hypothesis
64 keras
65 pytestCheckHook
66 tensorflow
67 torchmetrics
68 ];
69
70 pythonImportsCheck = [ "treex" ];
71
72 meta = with lib; {
73 description = "Pytree Module system for Deep Learning in JAX";
74 homepage = "https://github.com/cgarciae/treex";
75 license = licenses.mit;
76 maintainers = with maintainers; [ ndl ];
77 };
78}