diff --git a/pkgs/development/python-modules/treex/default.nix b/pkgs/development/python-modules/treex/default.nix
index 1f4a55416d56..0b5ad0c89839 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
@@ -38,6 +39,8 @@ buildPythonPackage rec {
     poetry-core
   ];
 
+  buildInputs = [ jaxlib ];
+
   propagatedBuildInputs = [
     einops
     flax