python3Packages.elegy: init at 0.8.4

3 files changed, 86 insertions, 3 deletions

diff --git a/pkgs/development/python-modules/einops/default.nix b/pkgs/development/python-modules/einops/default.nix
index a2de3a70971d5..18b8d671c98af 100644
--- a/pkgs/development/python-modules/einops/default.nix
+++ b/pkgs/development/python-modules/einops/default.nix
@@ -44,6 +44,11 @@ buildPythonPackage rec {
 
   checkPhase = ''
     export HOME=$TMPDIR
+
+    # Prevent hangs on PyTorch-related tests, see
+    # https://discuss.pytorch.org/t/pytorch-cpu-hangs-on-nn-linear/17748/4
+    export OMP_NUM_THREADS=1
+
     nosetests -v -w tests
   '';
 
diff --git a/pkgs/development/python-modules/elegy/default.nix b/pkgs/development/python-modules/elegy/default.nix
new file mode 100644
index 0000000000000..ec968dce8ec3a
--- /dev/null
+++ b/pkgs/development/python-modules/elegy/default.nix
@@ -0,0 +1,75 @@
+{ buildPythonPackage
+, cloudpickle
+, deepdish
+, deepmerge
+, dm-haiku
+, fetchFromGitHub
+, lib
+, poetry
+, pytestCheckHook
+, pytorch
+, pyyaml
+, sh
+, tables
+, tabulate
+, tensorboardx
+, tensorflow
+, toolz
+, treex
+, typing-extensions
+}:
+
+buildPythonPackage rec {
+  pname = "elegy";
+  version = "0.8.4";
+  format = "pyproject";
+
+  src = fetchFromGitHub {
+    owner = "poets-ai";
+    repo = pname;
+    rev = version;
+    sha256 = "11w8lgl31b52w2qri8j8cgzd30sn8i3769g8nkkshvgkjgca9r4g";
+  };
+
+  nativeBuildInputs = [
+    poetry
+  ];
+
+  propagatedBuildInputs = [
+    cloudpickle
+    deepdish
+    deepmerge
+    dm-haiku
+    pyyaml
+    tables
+    tabulate
+    tensorboardx
+    toolz
+    treex
+    typing-extensions
+  ];
+
+  pythonImportsCheck = [
+    "elegy"
+  ];
+
+  checkInputs = [
+    pytestCheckHook
+    pytorch
+    sh
+    tensorflow
+  ];
+
+  disabledTests = [
+    # Fails with `Could not find compiler for platform Host: NOT_FOUND: could not find registered compiler for platform Host -- check target linkage`.
+    # Runs fine in docker with Ubuntu 22.04. I suspect the issue is the sandboxing in `nixpkgs` but not sure.
+    "test_saved_model_poly"
+  ];
+
+  meta = with lib; {
+    description = "Neural Networks framework based on Jax inspired by Keras and Haiku";
+    homepage = "https://github.com/poets-ai/elegy";
+    license = licenses.asl20;
+    maintainers = with maintainers; [ ndl ];
+  };
+}
diff --git a/pkgs/development/python-modules/treex/relax-deps.patch b/pkgs/development/python-modules/treex/relax-deps.patch
index 5b00082ba30c8..347098a4a6d1b 100644
--- a/pkgs/development/python-modules/treex/relax-deps.patch
+++ b/pkgs/development/python-modules/treex/relax-deps.patch
@@ -1,10 +1,13 @@
 diff --git a/pyproject.toml b/pyproject.toml
-index 7b9ef68..ec11f32 100644
+index f0ff8a0..56787ca 100644
 --- a/pyproject.toml
 +++ b/pyproject.toml
-@@ -18,7 +18,7 @@ python = "^3.7"
+@@ -16,9 +16,9 @@ secondary = true
+ [tool.poetry.dependencies]
+ python = "^3.7"
  flax = "^0.3.4"
- PyYAML = "^5.4.1"
+-PyYAML = "^5.4.1"
++PyYAML = ">=5.4.1"
  rich = "^10.7.0"
 -optax = "^0.0.9"
 +optax = ">=0.0.9"