tjh.dev
1{ stdenv
2, fetchurl
3, buildPythonPackage
4, isPy36, isPy35, isPy27
5, cudaSupport ? false
6, cudatoolkit ? null
7, cudnn ? null
8, linuxPackages ? null
9, numpy
10, six
11, protobuf3_2
12, swig
13, werkzeug
14, mock
15, zlib
16}:
17
18assert cudaSupport -> cudatoolkit != null
19 && cudnn != null
20 && linuxPackages != null;
21
22# unsupported combination
23assert ! (stdenv.isDarwin && cudaSupport);
24
25# tensorflow is built from a downloaded wheel, because the upstream
26# project's build system is an arcane beast based on
27# bazel. Untangling it and building the wheel from source is an open
28# problem.
29
30buildPythonPackage rec {
31 pname = "tensorflow";
32 version = "1.1.0";
33 name = "${pname}-${version}";
34 format = "wheel";
35 disabled = ! (isPy36 || isPy35 || isPy27);
36
37 src = let
38 tfurl = sys: proc: pykind:
39 let
40 tfpref = if proc == "gpu"
41 then "gpu/tensorflow_gpu"
42 else "cpu/tensorflow";
43 in
44 "https://storage.googleapis.com/tensorflow/${sys}/${tfpref}-${version}-${pykind}.whl";
45 dls =
46 {
47 darwin.cpu = {
48 py2 = {
49 url = tfurl "mac" "cpu" "py2-none-any" ;
50 sha256 = "1fgf26lw0liqxc9pywc8y2mj8l1mv48nhkav0pag9vavdacb9mqr";
51 };
52 py3 = {
53 url = tfurl "mac" "cpu" "py3-none-any" ;
54 sha256 = "0z5p1fra7bih0vqn618i2w3vyy8d1rkc72k7bmjq0rw8msl717ia";
55 };
56 };
57 linux-x86_64.cpu = {
58 py2 = {
59 url = tfurl "linux" "cpu" "cp27-none-linux_x86_64";
60 sha256 = "0ld3hqx3idxk0zcrvn3p9yqnmx09zsj3mw66jlfw6fkv5hznx8j2";
61 };
62 py35 = {
63 url = tfurl "linux" "cpu" "cp35-cp35m-linux_x86_64";
64 sha256 = "0ahz9222rzqrk43lb9w4m351klkm6mlnnvw8xfqip28vbmymw90b";
65 };
66 py36 = {
67 url = tfurl "linux" "cpu" "cp36-cp36m-linux_x86_64";
68 sha256 = "1a2cc8ihl94iqff76nxg6bq85vfb7sj5cvvi9sxy2f43k32fi4lv";
69 };
70
71 };
72 linux-x86_64.cuda = {
73 py2 = {
74 url = tfurl "linux" "gpu" "cp27-none-linux_x86_64";
75 sha256 = "1baa9jwr6f8f62dyx6isbw8yyrd0pi1dz1srjblfqsyk1x3pnfvh";
76 };
77 py35 = {
78 url = tfurl "linux" "gpu" "cp35-cp35m-linux_x86_64";
79 sha256 = "0606m2awy0ifhniy8lsyhd0xc388dgrwksn87989xlgy90wpxi92";
80 };
81 py36 = {
82 url = tfurl "linux" "gpu" "cp36-cp36m-linux_x86_64";
83 sha256 = "0lvbmfa87qzrajadpsf13gi3l71vryzkryzqfvkykivqrdjsvj8q";
84 };
85
86 };
87 };
88 in
89 fetchurl (
90 if stdenv.isDarwin then
91 if isPy27 then
92 dls.darwin.cpu.py2
93 else
94 dls.darwin.cpu.py3
95 else if isPy36 then
96 if cudaSupport then
97 dls.linux-x86_64.cuda.py36
98 else dls.linux-x86_64.cpu.py36
99 else if isPy35 then
100 if cudaSupport then
101 dls.linux-x86_64.cuda.py35
102 else dls.linux-x86_64.cpu.py35
103 else
104 if cudaSupport then
105 dls.linux-x86_64.cuda.py2
106 else
107 dls.linux-x86_64.cpu.py2
108 );
109
110 propagatedBuildInputs = with stdenv.lib;
111 [ numpy six protobuf3_2 swig werkzeug mock ]
112 ++ optionals cudaSupport [ cudatoolkit cudnn stdenv.cc ];
113
114 # Note that we need to run *after* the fixup phase because the
115 # libraries are loaded at runtime. If we run in preFixup then
116 # patchelf --shrink-rpath will remove the cuda libraries.
117 postFixup = let
118 rpath = stdenv.lib.makeLibraryPath
119 (if cudaSupport then
120 [ stdenv.cc.cc.lib zlib cudatoolkit cudnn
121 linuxPackages.nvidia_x11 ]
122 else
123 [ stdenv.cc.cc.lib zlib ]
124 );
125 in
126 ''
127 find $out -name '*.so' -exec patchelf --set-rpath "${rpath}" {} \;
128 '';
129
130 doCheck = false;
131
132 meta = with stdenv.lib; {
133 description = "TensorFlow helps the tensors flow";
134 homepage = http://tensorflow.org;
135 license = licenses.asl20;
136 maintainers = with maintainers; [ jpbernardy ];
137 platforms = with platforms; if cudaSupport then linux else linux ++ darwin;
138 };
139}