diff options
Diffstat (limited to 'pkgs/development/libraries/libtensorflow')
-rw-r--r-- | pkgs/development/libraries/libtensorflow/default.nix | 57 |
1 files changed, 47 insertions, 10 deletions
diff --git a/pkgs/development/libraries/libtensorflow/default.nix b/pkgs/development/libraries/libtensorflow/default.nix index fce15b121b4..c8b299034fc 100644 --- a/pkgs/development/libraries/libtensorflow/default.nix +++ b/pkgs/development/libraries/libtensorflow/default.nix @@ -1,11 +1,50 @@ -{ stdenv, fetchurl, patchelf }: -stdenv.mkDerivation rec { +{ stdenv +, fetchurl +, patchelf +, cudaSupport ? false, symlinkJoin, cudatoolkit, cudnn, nvidia_x11 +}: +with stdenv.lib; +let + tfType = if cudaSupport then "gpu" else "cpu"; + system = + if stdenv.isx86_64 + then if stdenv.isLinux then "linux-x86_64" + else if stdenv.isDarwin then "darwin-x86_64" else unavailable + else unavailable; + unavailable = throw "libtensorflow is not available for this platform!"; + cudatoolkit_joined = symlinkJoin { + name = "unsplit_cudatoolkit"; + paths = [ cudatoolkit.out + cudatoolkit.lib ];}; + rpath = makeLibraryPath ([stdenv.cc.libc stdenv.cc.cc.lib] ++ + optionals cudaSupport [ cudatoolkit_joined cudnn nvidia_x11 ]); + patchLibs = + if stdenv.isDarwin + then '' + install_name_tool -id $out/lib/libtensorflow.so $out/lib/libtensorflow.so + install_name_tool -id $out/lib/libtensorflow_framework.so $out/lib/libtensorflow_framework.so + '' + else '' + ${patchelf}/bin/patchelf --set-rpath "${rpath}:$out/lib" $out/lib/libtensorflow.so + ${patchelf}/bin/patchelf --set-rpath "${rpath}" $out/lib/libtensorflow_framework.so + ''; + +in stdenv.mkDerivation rec { pname = "libtensorflow"; version = "1.8.0"; name = "${pname}-${version}"; src = fetchurl { - url = "https://storage.googleapis.com/tensorflow/${pname}/${pname}-cpu-linux-x86_64-${version}.tar.gz"; - sha256 = "0qzy15rc3x961cyi3bqnygrcnw4x69r28xkwhpwrv1r0gi6k73ha"; + url = "https://storage.googleapis.com/tensorflow/${pname}/${pname}-${tfType}-${system}-${version}.tar.gz"; + sha256 = + if system == "linux-x86_64" then + if cudaSupport + then "0m1g4sqr9as0jgfx7wlyay2nkad6wgvsyk2gvhfkqkq5sm1vbx85" + else "0qzy15rc3x961cyi3bqnygrcnw4x69r28xkwhpwrv1r0gi6k73ha" + else if system == "darwin-x86_64" then + if cudaSupport + then unavailable + else "0q8lmyj8l50hl6l48c640ixanvhqf2836bicyl9p2x8sj97b7y8l" + else unavailable; }; # Patch library to use our libc, libstdc++ and others @@ -15,18 +54,16 @@ stdenv.mkDerivation rec { tar -C $out -xzf $src chmod +w $out/lib/libtensorflow.so chmod +w $out/lib/libtensorflow_framework.so - ${patchelf}/bin/patchelf --set-rpath "${stdenv.cc.libc}/lib:${stdenv.cc.cc.lib}/lib:$out/lib" $out/lib/libtensorflow.so - ${patchelf}/bin/patchelf --set-rpath "${stdenv.cc.libc}/lib:${stdenv.cc.cc.lib}/lib" $out/lib/libtensorflow_framework.so + ${patchLibs} chmod -w $out/lib/libtensorflow.so chmod -w $out/lib/libtensorflow_framework.so ''; - meta = with stdenv.lib; { - inherit version; + meta = { description = "C API for TensorFlow"; + homepage = https://www.tensorflow.org/versions/master/install/install_c; license = licenses.asl20; + platforms = with platforms; linux ++ darwin; maintainers = [maintainers.basvandijk]; - platforms = platforms.linux; - homepage = https://www.tensorflow.org/versions/master/install/install_c; }; } |