@@ -0,0 +1,75 @@
1+{ buildPythonPackage
2+, cloudpickle
3+, deepdish
4+, deepmerge
5+, dm-haiku
6+, fetchFromGitHub
7+, lib
8+, poetry
9+, pytestCheckHook
10+, pytorch
11+, pyyaml
12+, sh
13+, tables
14+, tabulate
15+, tensorboardx
16+, tensorflow
17+, toolz
18+, treex
19+, typing-extensions
20+}:
21+
22+buildPythonPackage rec {
23+ pname = "elegy";
24+ version = "0.8.4";
25+ format = "pyproject";
26+
27+ src = fetchFromGitHub {
28+ owner = "poets-ai";
29+ repo = pname;
30+ rev = version;
31+ sha256 = "11w8lgl31b52w2qri8j8cgzd30sn8i3769g8nkkshvgkjgca9r4g";
32+ };
33+
34+ nativeBuildInputs = [
35+ poetry
36+ ];
37+
38+ propagatedBuildInputs = [
39+ cloudpickle
40+ deepdish
41+ deepmerge
42+ dm-haiku
43+ pyyaml
44+ tables
45+ tabulate
46+ tensorboardx
47+ toolz
48+ treex
49+ typing-extensions
50+ ];
51+
52+ pythonImportsCheck = [
53+ "elegy"
54+ ];
55+
56+ checkInputs = [
57+ pytestCheckHook
58+ pytorch
59+ sh
60+ tensorflow
61+ ];
62+
63+ disabledTests = [
64+ # Fails with `Could not find compiler for platform Host: NOT_FOUND: could not find registered compiler for platform Host -- check target linkage`.
65+ # Runs fine in docker with Ubuntu 22.04. I suspect the issue is the sandboxing in `nixpkgs` but not sure.
66+ "test_saved_model_poly"
67+ ];
68+
69+ meta = with lib; {
70+ description = "Neural Networks framework based on Jax inspired by Keras and Haiku";
71+ homepage = "https://github.com/poets-ai/elegy";
72+ license = licenses.asl20;
73+ maintainers = with maintainers; [ ndl ];
74+ };
75+}