summary refs log tree commit diff
path: root/pkgs/development/python-modules/tensorflow
diff options
context:
space:
mode:
authorJason Miller <jcmiller@pdtpartners.com>2019-09-04 15:44:24 -0400
committerJohn Ericson <John.Ericson@Obsidian.Systems>2020-03-31 12:01:22 -0400
commit4a4b44803cdc5c1bd4458d77aec71cb9e374d837 (patch)
tree3934079a4db169aefbda927dd38b1fece0e12a4d /pkgs/development/python-modules/tensorflow
parentcef0110b8452e236c1762f5a9e348f29ad48ff84 (diff)
downloadnixpkgs-4a4b44803cdc5c1bd4458d77aec71cb9e374d837.tar
nixpkgs-4a4b44803cdc5c1bd4458d77aec71cb9e374d837.tar.gz
nixpkgs-4a4b44803cdc5c1bd4458d77aec71cb9e374d837.tar.bz2
nixpkgs-4a4b44803cdc5c1bd4458d77aec71cb9e374d837.tar.lz
nixpkgs-4a4b44803cdc5c1bd4458d77aec71cb9e374d837.tar.xz
nixpkgs-4a4b44803cdc5c1bd4458d77aec71cb9e374d837.tar.zst
nixpkgs-4a4b44803cdc5c1bd4458d77aec71cb9e374d837.zip
tensorflow: bugfix for bazel settings (intel mkl, dnnl)
Diffstat (limited to 'pkgs/development/python-modules/tensorflow')
-rw-r--r--pkgs/development/python-modules/tensorflow/1/default.nix8
-rw-r--r--pkgs/development/python-modules/tensorflow/2/default.nix8
2 files changed, 14 insertions, 2 deletions
diff --git a/pkgs/development/python-modules/tensorflow/1/default.nix b/pkgs/development/python-modules/tensorflow/1/default.nix
index 78843c612ee..1b7c338d2a1 100644
--- a/pkgs/development/python-modules/tensorflow/1/default.nix
+++ b/pkgs/development/python-modules/tensorflow/1/default.nix
@@ -18,6 +18,7 @@
 # it would also make the default tensorflow package unfree. See
 # https://groups.google.com/a/tensorflow.org/forum/#!topic/developers/iRCt5m4qUz0
 , cudaSupport ? false, nvidia_x11 ? null, cudatoolkit ? null, cudnn ? null, nccl ? null
+, mklSupport ? false, mkl ? null
 # XLA without CUDA is broken
 , xlaSupport ? cudaSupport
 # Default from ./configure script
@@ -36,6 +37,8 @@ assert cudaSupport -> nvidia_x11 != null
 # unsupported combination
 assert ! (stdenv.isDarwin && cudaSupport);
 
+assert mklSupport -> mkl != null;
+
 let
   withTensorboard = pythonOlder "3.6";
 
@@ -167,6 +170,8 @@ let
       cudatoolkit
       cudnn
       nvidia_x11
+    ] ++ lib.optionals mklSupport [
+      mkl
     ] ++ lib.optionals stdenv.isDarwin [
       Foundation
       Security
@@ -284,7 +289,8 @@ let
     ];
     bazelBuildFlags = [
       "--config=opt" # optimize using the flags set in the configure phase
-    ];
+    ]
+    ++ lib.optionals (mklSupport) [ "--config=mkl" ];
 
     bazelTarget = "//tensorflow/tools/pip_package:build_pip_package //tensorflow/tools/lib_package:libtensorflow";
 
diff --git a/pkgs/development/python-modules/tensorflow/2/default.nix b/pkgs/development/python-modules/tensorflow/2/default.nix
index 29430323084..2867308f9c4 100644
--- a/pkgs/development/python-modules/tensorflow/2/default.nix
+++ b/pkgs/development/python-modules/tensorflow/2/default.nix
@@ -18,6 +18,7 @@
 # it would also make the default tensorflow package unfree. See
 # https://groups.google.com/a/tensorflow.org/forum/#!topic/developers/iRCt5m4qUz0
 , cudaSupport ? false, nvidia_x11 ? null, cudatoolkit ? null, cudnn ? null, nccl ? null
+, mklSupport ? false, mkl ? null
 # XLA without CUDA is broken
 , xlaSupport ? cudaSupport
 # Default from ./configure script
@@ -36,6 +37,8 @@ assert cudaSupport -> nvidia_x11 != null
 # unsupported combination
 assert ! (stdenv.isDarwin && cudaSupport);
 
+assert mklSupport -> mkl != null;
+
 let
   withTensorboard = pythonOlder "3.6";
 
@@ -163,6 +166,8 @@ let
       cudatoolkit
       cudnn
       nvidia_x11
+    ] ++ lib.optionals mklSupport [
+      mkl
     ] ++ lib.optionals stdenv.isDarwin [
       Foundation
       Security
@@ -278,7 +283,8 @@ let
     ];
     bazelBuildFlags = [
       "--config=opt" # optimize using the flags set in the configure phase
-    ];
+    ]
+    ++ lib.optionals (mklSupport) [ "--config=mkl" ];
 
     bazelTarget = "//tensorflow/tools/pip_package:build_pip_package //tensorflow/tools/lib_package:libtensorflow";