1 files changed, 11 insertions, 5 deletions
diff --git a/pkgs/development/python-modules/treex/default.nix b/pkgs/development/python-modules/treex/default.nix
index bdea03896cb58..7ed83adc64d15 100644
--- a/pkgs/development/python-modules/treex/default.nix
+++ b/pkgs/development/python-modules/treex/default.nix
@@ -5,6 +5,7 @@
 , fetchFromGitHub
 , flax
 , hypothesis
+, jaxlib
 , keras
 , lib
 , poetry-core
@@ -17,24 +18,29 @@
 
 buildPythonPackage rec {
   pname = "treex";
-  version = "0.6.7";
+  version = "0.6.9";
   format = "pyproject";
 
   src = fetchFromGitHub {
     owner = "cgarciae";
     repo = pname;
     rev = version;
-    sha256 = "1hl3wj71c7cp7jzkhyjy7xgs2vc8c89icq0bgfr49y4pwv69n43m";
+    sha256 = "1yvlldmhji12h249j14ba44hnb9x1fhrj7rh1cx2vn0vxj5wpg7x";
   };
 
-  patches = [
-    ./relax-deps.patch
-  ];
+  postPatch = ''
+    substituteInPlace pyproject.toml \
+      --replace 'rich = "^10.7.0"' 'rich = ">=10.7.0"' \
+      --replace 'PyYAML = "^5.4.1"' 'PyYAML = ">=5.4.1"' \
+      --replace 'optax = "^0.0.9"' 'optax = ">=0.0.9"'
+  '';
 
   nativeBuildInputs = [
     poetry-core
   ];
 
+  buildInputs = [ jaxlib ];
+
   propagatedBuildInputs = [
     einops
     flax