tjh.dev
1{ stdenv
2, lib
3, fetchurl
4, buildPythonPackage
5, isPy3k, isPy36, pythonOlder
6, astor
7, gast
8, numpy
9, six
10, termcolor
11, protobuf
12, absl-py
13, grpcio
14, mock
15, backports_weakref
16, enum34
17, tensorflow-estimator
18, tensorflow-tensorboard
19, cudaSupport ? false
20, cudatoolkit ? null
21, cudnn ? null
22, nvidia_x11 ? null
23, zlib
24, python
25, symlinkJoin
26, keras-applications
27, keras-preprocessing
28}:
29
30# We keep this binary build for two reasons:
31# - the source build doesn't work on Darwin.
32# - the source build is currently brittle and not easy to maintain
33
34assert cudaSupport -> cudatoolkit != null
35 && cudnn != null
36 && nvidia_x11 != null;
37let
38 cudatoolkit_joined = symlinkJoin {
39 name = "unsplit_cudatoolkit";
40 paths = [ cudatoolkit.out
41 cudatoolkit.lib ];};
42
43in buildPythonPackage rec {
44 pname = "tensorflow";
45 version = "1.13.1";
46 format = "wheel";
47
48 src = let
49 pyVerNoDot = lib.strings.stringAsChars (x: if x == "." then "" else x) "${python.pythonVersion}";
50 pyver = if stdenv.isDarwin then builtins.substring 0 1 pyVerNoDot else pyVerNoDot;
51 platform = if stdenv.isDarwin then "mac" else "linux";
52 unit = if cudaSupport then "gpu" else "cpu";
53 key = "${platform}_py_${pyver}_${unit}";
54 dls = import (./. + "/tf${version}-hashes.nix");
55 in fetchurl dls.${key};
56
57 propagatedBuildInputs = [ protobuf numpy termcolor grpcio six astor absl-py gast tensorflow-estimator tensorflow-tensorboard keras-applications keras-preprocessing ]
58 ++ lib.optional (!isPy3k) mock
59 ++ lib.optionals (pythonOlder "3.4") [ backports_weakref ];
60
61 # Upstream has a pip hack that results in bin/tensorboard being in both tensorflow
62 # and the propageted input tensorflow-tensorboard which causes environment collisions.
63 # another possibility would be to have tensorboard only in the buildInputs
64 # https://github.com/tensorflow/tensorflow/blob/v1.7.1/tensorflow/tools/pip_package/setup.py#L79
65 postInstall = ''
66 rm $out/bin/tensorboard
67 '';
68
69 # Note that we need to run *after* the fixup phase because the
70 # libraries are loaded at runtime. If we run in preFixup then
71 # patchelf --shrink-rpath will remove the cuda libraries.
72 postFixup = let
73 rpath = stdenv.lib.makeLibraryPath
74 ([ stdenv.cc.cc.lib zlib ] ++ lib.optionals cudaSupport [ cudatoolkit_joined cudnn nvidia_x11 ]);
75 in
76 lib.optionalString (stdenv.isLinux) ''
77 rrPath="$out/${python.sitePackages}/tensorflow/:$out/${python.sitePackages}/tensorflow/contrib/tensor_forest/:${rpath}"
78 internalLibPath="$out/${python.sitePackages}/tensorflow/python/_pywrap_tensorflow_internal.so"
79 find $out -name '*${stdenv.hostPlatform.extensions.sharedLibrary}' -exec patchelf --set-rpath "$rrPath" {} \;
80 '';
81
82
83 meta = with stdenv.lib; {
84 description = "Computation using data flow graphs for scalable machine learning";
85 homepage = http://tensorflow.org;
86 license = licenses.asl20;
87 maintainers = with maintainers; [ jyp abbradar ];
88 platforms = with platforms; linux ++ lib.optionals (!cudaSupport) darwin;
89 # Python 2.7 build uses different string encoding.
90 # See https://github.com/NixOS/nixpkgs/pull/37044#issuecomment-373452253
91 broken = stdenv.isDarwin && !isPy3k;
92 };
93}