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