pyrox.dev
1{ lib
2, buildPythonPackage
3, cloudpickle
4, deepdish
5, deepmerge
6, dm-haiku
7, fetchFromGitHub
8, jaxlib
9, poetry
10, pytestCheckHook
11, pythonOlder
12, pyyaml
13, sh
14, tables
15, tabulate
16, tensorboardx
17, tensorflow
18, toolz
19, torch
20, treex
21, typing-extensions
22}:
23
24buildPythonPackage rec {
25 pname = "elegy";
26 version = "0.8.6";
27 format = "pyproject";
28
29 disabled = pythonOlder "3.7";
30
31 src = fetchFromGitHub {
32 owner = "poets-ai";
33 repo = pname;
34 rev = version;
35 hash = "sha256-FZmLriYhsX+zyQKCtCjbOy6MH+AvjzHRNUyaDSXGlLI=";
36 };
37
38 # The cloudpickle constraint is too strict. wandb is marked as an optional
39 # dependency but `buildPythonPackage` doesn't seem to respect that setting.
40 # Python constraint: https://github.com/poets-ai/elegy/issues/244
41 postPatch = ''
42 substituteInPlace pyproject.toml \
43 --replace 'python = ">=3.7,<3.10"' 'python = ">=3.7"' \
44 --replace 'cloudpickle = "^1.5.0"' 'cloudpickle = "*"' \
45 --replace 'wandb = { version = "^0.12.10", optional = true }' ""
46 '';
47
48 nativeBuildInputs = [
49 poetry
50 ];
51
52 buildInputs = [
53 jaxlib
54 ];
55
56 propagatedBuildInputs = [
57 cloudpickle
58 deepdish
59 deepmerge
60 dm-haiku
61 pyyaml
62 tables
63 tabulate
64 tensorboardx
65 toolz
66 treex
67 typing-extensions
68 ];
69
70 pythonImportsCheck = [
71 "elegy"
72 ];
73
74 checkInputs = [
75 pytestCheckHook
76 torch
77 sh
78 tensorflow
79 ];
80
81 disabledTests = [
82 # Fails with `Could not find compiler for platform Host: NOT_FOUND: could not find registered compiler for platform Host -- check target linkage`.
83 # Runs fine in docker with Ubuntu 22.04. I suspect the issue is the sandboxing in `nixpkgs` but not sure.
84 "test_saved_model_poly"
85 # AttributeError: module 'jax' has no attribute 'tree_multimap'
86 "DataLoaderTestCase"
87 ];
88
89 meta = with lib; {
90 description = "Neural Networks framework based on Jax inspired by Keras and Haiku";
91 homepage = "https://github.com/poets-ai/elegy";
92 license = licenses.asl20;
93 maintainers = with maintainers; [ ndl ];
94 };
95}