Merge pull request #156808 from samuela/samuela/fix-jaxlib

JAX ecosystem: `jaxlib` should never be in `propagatedBuildInputs`
NixOS · Jan 28, 2022 · d8ff157 · d8ff157
2 parents 65a33fe + 93f88be
commit d8ff157
Show file tree

Hide file tree

Showing 7 changed files with 17 additions and 5 deletions.
diff --git a/pkgs/development/python-modules/dm-haiku/default.nix b/pkgs/development/python-modules/dm-haiku/default.nix
@@ -4,6 +4,7 @@
 , dill
 , dm-tree
 , fetchFromGitHub
+, jaxlib
 , jmp
 , lib
 , pytestCheckHook
@@ -31,6 +32,7 @@ buildPythonPackage rec {
     chex
     cloudpickle
     dm-tree
+    jaxlib
     pytestCheckHook
     tensorflow
   ];

diff --git 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

diff --git 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

diff --git a/pkgs/development/python-modules/jmp/default.nix b/pkgs/development/python-modules/jmp/default.nix
@@ -19,17 +19,17 @@ buildPythonPackage rec {
     sha256 = "0hh4cmp93wjyidj48gh07vhx2kjvpwd23xvy79bsjn5qaaf6q4cm";
   };
 
-  # Wheel requires only `numpy`, but the import needs both `jax` and `jaxlib`.
+  # Wheel requires only `numpy`, but the import needs `jax`.
   propagatedBuildInputs = [
     jax
-    jaxlib
   ];
 
   pythonImportsCheck = [
     "jmp"
   ];
 
   checkInputs = [
+    jaxlib
     pytestCheckHook
   ];
 

diff --git a/pkgs/development/python-modules/optax/default.nix b/pkgs/development/python-modules/optax/default.nix
@@ -22,10 +22,11 @@ buildPythonPackage rec {
     sha256 = "1q8cxc42a5xais2ll1l238cnn3l7w28savhgiz0lg01ilz2ysbli";
   };
 
+  buildInputs = [ jaxlib ];
+
   propagatedBuildInputs = [
     absl-py
     chex
-    jaxlib
     numpy
   ];
 

diff --git a/pkgs/development/python-modules/treeo/default.nix b/pkgs/development/python-modules/treeo/default.nix
@@ -22,12 +22,12 @@ buildPythonPackage rec {
     poetry-core
   ];
 
-  # These deps are not needed for the wheel, but required during the import.
+  # jax is not declared in the dependencies, but is necessary.
   propagatedBuildInputs = [
     jax
-    jaxlib
   ];
 
+  checkInputs = [ jaxlib ];
   pythonImportsCheck = [
     "treeo"
   ];

diff --git 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
-Original file line number
+Diff line change
@@ Expand Up / @@ -22,10 +22,11 @@ buildPythonPackage rec { @@
         sha256 = "1q8cxc42a5xais2ll1l238cnn3l7w28savhgiz0lg01ilz2ysbli";
       };
+      buildInputs = [ jaxlib ];
       propagatedBuildInputs = [
         absl-py
         chex
-        jaxlib
         numpy
       ];
@@ Expand Down @@