Merge pull request #23240 from jyp/tf

tensorflow: depend on cudatoolkit 8 and cudnn 5.1

pkgs/development/python-modules/tensorflow/default.nix

@@ -3,8 +3,8 @@
 , buildPythonPackage
 , isPy35, isPy27
 , cudaSupport ? false
-, cudatoolkit75 ? null
-, cudnn5_cudatoolkit75 ? null
+, cudatoolkit ? null
+, cudnn ? null
 , gcc49 ? null
 , linuxPackages ? null
 , numpy
@@ -16,8 +16,8 @@
 , zlib
 }:
 
-assert cudaSupport -> cudatoolkit75 != null
-                   && cudnn5_cudatoolkit75 != null
+assert cudaSupport -> cudatoolkit != null
+                   && cudnn != null
                    && gcc49 != null
                    && linuxPackages != null;
 
@@ -97,7 +97,7 @@ buildPythonPackage rec {
 
   propagatedBuildInputs = with stdenv.lib;
     [ numpy six protobuf3_2 swig mock ]
-    ++ optionals cudaSupport [ cudatoolkit75 cudnn5_cudatoolkit75 gcc49 ];
+    ++ optionals cudaSupport [ cudatoolkit cudnn gcc49 ];
 
   # Note that we need to run *after* the fixup phase because the
   # libraries are loaded at runtime. If we run in preFixup then
@@ -105,7 +105,7 @@ buildPythonPackage rec {
   postFixup = let
     rpath = stdenv.lib.makeLibraryPath
       (if cudaSupport then
-        [ gcc49.cc.lib zlib cudatoolkit75 cudnn5_cudatoolkit75
+        [ gcc49.cc.lib zlib cudatoolkit cudnn
           linuxPackages.nvidia_x11 ]
       else
         [ gcc.cc.lib zlib ]

pkgs/top-level/python-packages.nix

@@ -31278,6 +31278,8 @@ EOF
 
   tensorflowWithCuda = callPackage ../development/python-modules/tensorflow {
     cudaSupport = true;
+    cudatoolkit = pkgs.cudatoolkit8;
+    cudnn = pkgs.cudnn51_cudatoolkit80;
   };
 
   tflearn = buildPythonPackage rec {