tjh.dev
nixpkgs mirror (for testing)
github.com/NixOS/nixpkgs
nix
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}:
18
19buildPythonPackage rec {
20 pname = "treex";
21 version = "0.6.10";
22 format = "pyproject";
23
24 src = fetchFromGitHub {
25 owner = "cgarciae";
26 repo = pname;
27 rev = version;
28 hash = "sha256-ZHfgmRNbFh8DFZkmilY0pmRNQhJFqT689I7Lu8FuFm4=";
29 };
30
31 # At the time of writing (2022-03-29), rich is currently at version 11.0.0.
32 # The treeo dependency is compatible with a patch, but not marked as such in
33 # treex. See https://github.com/cgarciae/treex/issues/68.
34 postPatch = ''
35 substituteInPlace pyproject.toml \
36 --replace 'rich = "^11.2.0"' 'rich = "*"' \
37 --replace 'treeo = "^0.0.10"' 'treeo = "*"'
38 '';
39
40 nativeBuildInputs = [
41 poetry-core
42 ];
43
44 buildInputs = [ jaxlib ];
45
46 propagatedBuildInputs = [
47 einops
48 flax
49 pyyaml
50 rich
51 treeo
52 ];
53
54 checkInputs = [
55 cloudpickle
56 dm-haiku
57 hypothesis
58 keras
59 pytestCheckHook
60 tensorflow
61 ];
62
63 pythonImportsCheck = [
64 "treex"
65 ];
66
67 disabledTestPaths = [
68 # Require `torchmetrics` which is not packaged in `nixpkgs`.
69 "tests/metrics/test_mean_absolute_error.py"
70 "tests/metrics/test_mean_square_error.py"
71 ];
72
73 meta = with lib; {
74 description = "Pytree Module system for Deep Learning in JAX";
75 homepage = "https://github.com/cgarciae/treex";
76 license = licenses.mit;
77 maintainers = with maintainers; [ ndl ];
78 };
79}