diff options
| author | John Ericson <git@JohnEricson.me> | 2020-03-31 12:47:29 -0400 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2020-03-31 12:47:29 -0400 |
| commit | be0d1dc5a6c42456dbddde9b347d49cfaad034ef (patch) | |
| tree | 73b031513b6852871dda722cb57d03c6883f60a8 /pkgs | |
| parent | a90b82dfe9e89e579473591454f0f1c28b264bd5 (diff) | |
| parent | 4a4b44803cdc5c1bd4458d77aec71cb9e374d837 (diff) | |
Merge pull request #69454 from jmillerpdt/bugfix/tensorflow-mkl
pythonPackages.tensorflow: fix for bazel settings for intel mkl, dnnl
Diffstat (limited to 'pkgs')
| -rw-r--r-- | pkgs/development/python-modules/tensorflow/1/default.nix | 8 | ||||
| -rw-r--r-- | pkgs/development/python-modules/tensorflow/2/default.nix | 8 |
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 78843c612ee1..1b7c338d2a12 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 294303230847..2867308f9c43 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"; |