Merge pull request #69454 from jmillerpdt/bugfix/tensorflow-mkl

pythonPackages.tensorflow: fix for bazel settings for intel mkl, dnnl
This commit is contained in:
John Ericson 2020-03-31 12:47:29 -04:00 committed by GitHub
commit be0d1dc5a6
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 14 additions and 2 deletions

View File

@ -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";

View File

@ -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";