From f37b34fe7ced3c7a865a06668d7dbc45c85fe9bf Mon Sep 17 00:00:00 2001
From: Jean-Philippe Bernardy <jeanphilippe.bernardy@gmail.com>
Date: Mon, 27 Feb 2017 15:23:23 +0100
Subject: [PATCH] tensorflow: depend on cudatoolkit 8 and cudnn 5.1

---
 .../python-modules/tensorflow/default.nix            | 12 ++++++------
 pkgs/top-level/python-packages.nix                   |  2 ++
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/pkgs/development/python-modules/tensorflow/default.nix b/pkgs/development/python-modules/tensorflow/default.nix
index 9357ea90b40..04e820307a4 100644
--- a/pkgs/development/python-modules/tensorflow/default.nix
+++ b/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 ]
diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix
index 6a51e40d123..a268c811c52 100644
--- a/pkgs/top-level/python-packages.nix
+++ b/pkgs/top-level/python-packages.nix
@@ -31256,6 +31256,8 @@ EOF
 
   tensorflowWithCuda = callPackage ../development/python-modules/tensorflow {
     cudaSupport = true;
+    cudatoolkit = pkgs.cudatoolkit8;
+    cudnn = pkgs.cudnn51_cudatoolkit80;
   };
 
   tflearn = buildPythonPackage rec {