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