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