1 files changed, 5 insertions, 2 deletions
diff --git a/pkgs/development/python-modules/flax/default.nix b/pkgs/development/python-modules/flax/default.nix
index b8479c0f73cad..3ac8166bbbdd6 100644
--- a/pkgs/development/python-modules/flax/default.nix
+++ b/pkgs/development/python-modules/flax/default.nix
@@ -1,5 +1,6 @@
 { buildPythonPackage
 , fetchFromGitHub
+, jaxlib
 , keras
 , lib
 , matplotlib
@@ -12,15 +13,17 @@
 
 buildPythonPackage rec {
   pname = "flax";
-  version = "0.3.6";
+  version = "0.4.0";
 
   src = fetchFromGitHub {
     owner = "google";
     repo = pname;
     rev = "v${version}";
-    sha256 = "0zvq0vl88hiwmss49bnm7gdmndr1dfza2bcs1fj88a9r7w9dmlsr";
+    sha256 = "0rvdaxyf68qmm5d77gbizpcibyz2ic2pb2x7rgf7p8qwijyc39ws";
   };
 
+  buildInputs = [ jaxlib ];
+
   propagatedBuildInputs = [
     matplotlib
     msgpack