From c674a64ce671d3b8df1b6afa57e581dc47d793cd Mon Sep 17 00:00:00 2001
From: Samuel Ainsworth <skainsworth@gmail.com>
Date: Thu, 27 Jan 2022 20:52:07 +0000
Subject: [PATCH] python3Packages.flax: jaxlib belongs in buildInputs

---
 pkgs/development/python-modules/flax/default.nix | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/pkgs/development/python-modules/flax/default.nix b/pkgs/development/python-modules/flax/default.nix
index b8479c0f73ca..b5a22858b57a 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
@@ -21,6 +22,8 @@ buildPythonPackage rec {
     sha256 = "0zvq0vl88hiwmss49bnm7gdmndr1dfza2bcs1fj88a9r7w9dmlsr";
   };
 
+  buildInputs = [ jaxlib ];
+
   propagatedBuildInputs = [
     matplotlib
     msgpack