diff --git a/pkgs/development/python-modules/elegy/default.nix b/pkgs/development/python-modules/elegy/default.nix
index ec968dce8ec3..5b0cb293c0a3 100644
--- a/pkgs/development/python-modules/elegy/default.nix
+++ b/pkgs/development/python-modules/elegy/default.nix
@@ -4,6 +4,7 @@
 , deepmerge
 , dm-haiku
 , fetchFromGitHub
+, jaxlib
 , lib
 , poetry
 , pytestCheckHook
@@ -35,6 +36,8 @@ buildPythonPackage rec {
     poetry
   ];
 
+  buildInputs = [ jaxlib ];
+
   propagatedBuildInputs = [
     cloudpickle
     deepdish