Merge pull request #168707 from samuela/samuela/jax

python3Packages.jax: fix build when testing against jaxlibWithCuda
2022-04-17 12:31:12 -07:00
parent ba88a5afa6 cb7b514703
commit 06d2cfa7b4
1 changed files with 7 additions and 0 deletions
@@ -3,6 +3,7 @@
 , blas
 , buildPythonPackage
 , fetchFromGitHub
+, fetchpatch
 , jaxlib
 , lapack
 , numpy
@@ -34,6 +35,12 @@ buildPythonPackage rec {
  patches = [
    # See https://github.com/google/jax/issues/7944
    ./cache-fix.patch
+
+    # See https://github.com/google/jax/issues/10292
+    (fetchpatch {
+      url = "https://github.com/google/jax/commit/cadc8046d56e0c1433cf48a2f106947d5f4ecbfd.patch";
+      hash = "sha256-jrpIqt4LzWAswt/Cpwtfa5d1Yn31HcXkVH3ETmaigA0=";
+    })
  ];

  # jaxlib is _not_ included in propagatedBuildInputs because there are