python3Packages.jaxlib: unify fetchAttrs across all platforms

Based on code from NixOS#221390 by @uri-canva.
samuela · Mar 30, 2023 · aaf0d0a · aaf0d0a
1 parent d70820e
commit aaf0d0a
Showing 1 changed file with 10 additions and 7 deletions.
diff --git a/pkgs/development/python-modules/jaxlib/default.nix b/pkgs/development/python-modules/jaxlib/default.nix
@@ -232,14 +232,17 @@ let
       "--cxxopt=-x" "--cxxopt=c++" "--host_cxxopt=-x" "--host_cxxopt=c++"
     ];
 
+    # We intentionally overfetch so we can share the fetch derivation across all the different configurations.
     fetchAttrs = {
-      sha256 =
-        if cudaSupport then
-          "sha256-KN1akhisazX6bSAmCgPmw926c8lv232fMtgiouuwLb4="
-        else {
-          x86_64-linux = "sha256-3dYXATzG6gptSFpiPRa25668KRHYf5Z7SO5YTKsTUOI=";
-          aarch64-darwin = "sha256-QVAvjFe49wz6Jp43N8gWoo1PzQYTRSiFJ4R4CuEvyFw=";
-        }.${stdenv.system} or (throw "unsupported system ${stdenv.system}");
+      bazelFlags = [
+        "--config=avx_posix"
+        "--config=cuda"
+        "--config=mkl_open_source_only"
+      ];
+      preBuild = ''
+        export TF_SYSTEM_LIBS=""
+      '';
+      sha256 = "sha256-jYgvy8Ci+K8DtrnEvkLUbqxZbF5+bgC7RqkBkkOuCU4=";
     };
 
     buildAttrs = {