pythonPackages.tensorflow: 0.10.0 -> 1.0.0

2 files changed, 94 insertions, 68 deletions

diff --git a/pkgs/development/python-modules/tensorflow/cuda.nix b/pkgs/development/python-modules/tensorflow/cuda.nix
deleted file mode 100644
index 05a4cc3e4c18..000000000000
--- a/pkgs/development/python-modules/tensorflow/cuda.nix
+++ /dev/null
@@ -1,52 +0,0 @@
-{ stdenv
-, fetchurl
-, buildPythonPackage
-, swig
-, numpy
-, six
-, protobuf3_0
-, cudatoolkit75
-, cudnn5_cudatoolkit75
-, gcc49
-, zlib
-, linuxPackages
-, mock
-}:
-
-buildPythonPackage rec {
-  pname = "tensorflow";
-  version = "0.11.0rc0";
-  name = "${pname}-${version}";
-  format = "wheel";
-
-  src = fetchurl {
-    url = "https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow-${version}-cp27-none-linux_x86_64.whl";
-    sha256 = "1r8zlz95sw7bnjzg5zdbpa9dj8wmp8cvvgyl9sv3amsscagnnfj5";
-  };
-
-  buildInputs = [ swig ];
-  propagatedBuildInputs = [ numpy six protobuf3_0 cudatoolkit75 cudnn5_cudatoolkit75 gcc49 mock ];
-
-  # Note that we need to run *after* the fixup phase because the
-  # libraries are loaded at runtime. If we run in preFixup then
-  # patchelf --shrink-rpath will remove the cuda libraries.
-  postFixup = let
-    rpath = stdenv.lib.makeLibraryPath [
-      gcc49.cc.lib
-      zlib cudatoolkit75
-      cudnn5_cudatoolkit75
-      linuxPackages.nvidia_x11
-    ];
-  in ''
-    find $out -name '*.so' -exec patchelf --set-rpath "${rpath}" {} \;
-  '';
-
-  doCheck = false;
-
-  meta = with stdenv.lib; {
-    description = "TensorFlow helps the tensors flow (no gpu support)";
-    homepage = http://tensorflow.org;
-    license = licenses.asl20;
-    platforms   = platforms.linux;
-  };
-}
diff --git a/pkgs/development/python-modules/tensorflow/default.nix b/pkgs/development/python-modules/tensorflow/default.nix
index f8bc35eb5684..9357ea90b405 100644
--- a/pkgs/development/python-modules/tensorflow/default.nix
+++ b/pkgs/development/python-modules/tensorflow/default.nix
@@ -1,15 +1,29 @@
 { stdenv
 , fetchurl
 , buildPythonPackage
+, isPy35, isPy27
+, cudaSupport ? false
+, cudatoolkit75 ? null
+, cudnn5_cudatoolkit75 ? null
+, gcc49 ? null
+, linuxPackages ? null
 , numpy
 , six
-, protobuf3_0_0b2
+, protobuf3_2
 , swig
 , mock
 , gcc
 , zlib
 }:
 
+assert cudaSupport -> cudatoolkit75 != null
+                   && cudnn5_cudatoolkit75 != null
+                   && gcc49 != null
+                   && linuxPackages != null;
+
+# unsupported combination
+assert ! (stdenv.isDarwin && cudaSupport);
+
 # tensorflow is built from a downloaded wheel, because the upstream
 # project's build system is an arcane beast based on
 # bazel. Untangling it and building the wheel from source is an open
@@ -17,32 +31,96 @@
 
 buildPythonPackage rec {
   pname = "tensorflow";
-  version = "0.10.0";
+  version = "1.0.0";
   name = "${pname}-${version}";
   format = "wheel";
+  disabled = ! (isPy35 || isPy27);
 
-  src = fetchurl {
-    url = if stdenv.isDarwin then
-      "https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-${version}-py2-none-any.whl" else
-      "https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-${version}-cp27-none-linux_x86_64.whl";
-    sha256 = if stdenv.isDarwin then
-      "1gjybh3j3rn34bzhsxsfdbqgsr4jh50qyx2wqywvcb24fkvy40j9" else
-      "0g05pa4z6kdy0giz7hjgjgwf4zzr5l8cf1zh247ymixlikn3fnpx";
-  };
+  src = let
+      tfurl = sys: proc: pykind:
+        let
+          tfpref = if proc == "gpu"
+            then "gpu/tensorflow_gpu"
+            else "cpu/tensorflow";
+        in
+        "https://storage.googleapis.com/tensorflow/${sys}/${tfpref}-${version}-${pykind}.whl";
+      dls =
+        {
+        darwin.cpu = {
+          py2 = {
+            url = tfurl "mac" "cpu" "py2-none-any" ;
+            sha256 = "15ayil28p20wkgpwkr4mz0imjxnf049xx4117jspg1qkjg2bn1b2";
+          };
+          py3 = {
+            url = tfurl "mac" "cpu" "py3-none-any" ;
+            sha256 = "1ynyhbm7yrp421364s49a1r3p83zxy74iiy5c4hx2xm5c4gs29an";
+          };
+        };
+        linux-x86_64.cpu = {
+          py2 = {
+            url = tfurl "linux" "cpu" "cp27-none-linux_x86_64";
+            sha256 = "1hwhq1qhjrfkqfkxpsrq6mdmdibnqr3n7xvzkxp6gaqj73vn5ch2";
+          };
+          py3 = {
+            url = tfurl "linux" "cpu" "cp35-cp35m-linux_x86_64";
+            sha256 = "0jx2mmlw0nxah9l25r46i7diqiv31qcz7855n250lsxfwcppy7y3";
+          };
+        };
+        linux-x86_64.cuda = {
+          py2 = {
+            url = tfurl "linux" "gpu" "cp27-none-linux_x86_64";
+            sha256 = "0l8f71x3ama5a6idj05jrswlmp4yg37fxhz8lx2xmgk14aszbcy5";
+          };
+          py3 = {
+            url = tfurl "linux" "gpu" "cp35-cp35m-linux_x86_64";
+            sha256 = "12q7s0yk0h3r4glh0fhl1fcdx7jl8xikwwp04a1lcagasr51s36m";
+          };
+        };
+      };
+    in
+    fetchurl (
+      if stdenv.isDarwin then
+        if isPy35 then
+          dls.darwin.cpu.py3
+        else
+          dls.darwin.cpu.py2
+      else if isPy35 then
+        if cudaSupport then
+          dls.linux-x86_64.cuda.py3
+        else dls.linux-x86_64.cpu.py3
+      else
+        if cudaSupport then
+          dls.linux-x86_64.cuda.py2
+        else
+          dls.linux-x86_64.cpu.py2
+    );
 
-  propagatedBuildInputs = [ numpy six protobuf3_0_0b2 swig mock];
+  propagatedBuildInputs = with stdenv.lib;
+    [ numpy six protobuf3_2 swig mock ]
+    ++ optionals cudaSupport [ cudatoolkit75 cudnn5_cudatoolkit75 gcc49 ];
 
-  preFixup = ''
-    RPATH="${stdenv.lib.makeLibraryPath [ gcc.cc.lib zlib ]}"
-    find $out -name '*.so' -exec patchelf --set-rpath "$RPATH" {} \;
+  # Note that we need to run *after* the fixup phase because the
+  # libraries are loaded at runtime. If we run in preFixup then
+  # patchelf --shrink-rpath will remove the cuda libraries.
+  postFixup = let
+    rpath = stdenv.lib.makeLibraryPath
+      (if cudaSupport then
+        [ gcc49.cc.lib zlib cudatoolkit75 cudnn5_cudatoolkit75
+          linuxPackages.nvidia_x11 ]
+      else
+        [ gcc.cc.lib zlib ]
+      );
+  in
+  ''
+    find $out -name '*.so' -exec patchelf --set-rpath "${rpath}" {} \;
   '';
 
   doCheck = false;
 
   meta = with stdenv.lib; {
-    description = "TensorFlow helps the tensors flow (no gpu support)";
+    description = "TensorFlow helps the tensors flow";
     homepage = http://tensorflow.org;
     license = licenses.asl20;
-    platforms = with platforms; linux ++ darwin;
+    platforms = with platforms; if cudaSupport then linux else linux ++ darwin;
   };
 }