python3Packages.tensorflow_2: 2.1.0 -> 2.3.0

Also:
- patch to remove scipy requirement
- add cuda to RPATH
- don’t include nvidia_x11 (This isn’t needed, we can get it from
/run/opengl-driver being in the RPATH.)

Co-authored-by: Arnout Engelen <arnout@bzzt.net>
Co-authored-by: Daniël de Kok <me@github.danieldk.eu>

1 files changed, 40 insertions, 44 deletions

diff --git a/pkgs/development/python-modules/tensorflow/2/default.nix b/pkgs/development/python-modules/tensorflow/2/default.nix
index eedd6e6d0df..f99a2590c44 100644
--- a/pkgs/development/python-modules/tensorflow/2/default.nix
+++ b/pkgs/development/python-modules/tensorflow/2/default.nix
@@ -1,4 +1,4 @@
-{ stdenv, pkgs, bazel_0_29, buildBazelPackage, lib, fetchFromGitHub, fetchpatch, symlinkJoin
+{ stdenv, pkgs, bazel_3, buildBazelPackage, lib, fetchFromGitHub, fetchpatch, symlinkJoin
 , addOpenGLRunpath
 # Python deps
 , buildPythonPackage, isPy3k, isPy27, pythonOlder, pythonAtLeast, python
@@ -6,7 +6,7 @@
 , numpy, tensorflow-tensorboard_2, backports_weakref, mock, enum34, absl-py
 , future, setuptools, wheel, keras-preprocessing, keras-applications, google-pasta
 , functools32
-, opt-einsum
+, opt-einsum, astunparse, h5py
 , termcolor, grpcio, six, wrapt, protobuf, tensorflow-estimator_2
 # Common deps
 , git, swig, which, binutils, glibcLocales, cython
@@ -17,7 +17,7 @@
 # that in nix as well. It would make some things easier and less confusing, but
 # it would also make the default tensorflow package unfree. See
 # https://groups.google.com/a/tensorflow.org/forum/#!topic/developers/iRCt5m4qUz0
-, cudaSupport ? false, nvidia_x11 ? null, cudatoolkit ? null, cudnn ? null, nccl ? null
+, cudaSupport ? false, cudatoolkit ? null, cudnn ? null, nccl ? null
 , mklSupport ? false, mkl ? null
 # XLA without CUDA is broken
 , xlaSupport ? cudaSupport
@@ -30,8 +30,7 @@
 , Foundation, Security
 }:
 
-assert cudaSupport -> nvidia_x11 != null
-                   && cudatoolkit != null
+assert cudaSupport -> cudatoolkit != null
                    && cudnn != null;
 
 # unsupported combination
@@ -72,7 +71,7 @@ let
 
   tfFeature = x: if x then "1" else "0";
 
-  version = "2.1.0";
+  version = "2.3.0";
   variant = if cudaSupport then "-gpu" else "";
   pname = "tensorflow${variant}";
 
@@ -97,40 +96,32 @@ let
 
   bazel-build = buildBazelPackage {
     name = "${pname}-${version}";
-    bazel = bazel_0_29;
+    bazel = bazel_3;
 
     src = fetchFromGitHub {
       owner = "tensorflow";
       repo = "tensorflow";
       rev = "v${version}";
-      sha256 = "1g79xi8yl4sjia8ysk9b7xfzrz83zy28v5dlb2wzmcf0k5pmz60p";
+      sha256 = "1dd5fgyiazyfy7y2iv4v42qnap51fr6dzwb26inrsj7aaas06j71";
     };
 
     patches = [
-      # Work around https://github.com/tensorflow/tensorflow/issues/24752
-      ../no-saved-proto.patch
       # Fixes for NixOS jsoncpp
       ../system-jsoncpp.patch
 
+      ./lift-gast-restriction.patch
+
+      # see https://github.com/tensorflow/tensorflow/issues/40688
       (fetchpatch {
-        name = "backport-pr-18950.patch";
-        url = "https://github.com/tensorflow/tensorflow/commit/73640aaec2ab0234d9fff138e3c9833695570c0a.patch";
-        sha256 = "1n9ypbrx36fc1kc9cz5b3p9qhg15xxhq4nz6ap3hwqba535nakfz";
+        url = "https://github.com/tensorflow/tensorflow/commit/75ea0b31477d6ba9e990e296bbbd8ca4e7eebadf.patch";
+        sha256 = "1xp1icacig0xm0nmb05sbrf4nw4xbln9fhc308birrv8286zx7wv";
       })
 
+      # see https://github.com/tensorflow/tensorflow/issues/40884
       (fetchpatch {
-        # Don't try to fetch things that don't exist
-        name = "prune-missing-deps.patch";
-        url = "https://github.com/tensorflow/tensorflow/commit/b39b1ed24b4814db27d2f748dc85c10730ae851d.patch";
-        sha256 = "1skysz53nancvw1slij6s7flar2kv3gngnsq60ff4lap88kx5s6c";
-        excludes = [ "tensorflow/cc/saved_model/BUILD" ];
+        url = "https://github.com/tensorflow/tensorflow/pull/41867/commits/65341f73d110bf173325768947343e1bb8f699fc.patch";
+        sha256 = "18ykkycaag1pcarz53bz6ydxjlah92j4178qn58gcayx1fy7hvh3";
       })
-
-      ./lift-gast-restriction.patch
-
-      # cuda 10.2 does not have "-bin2c-path" option anymore
-      # https://github.com/tensorflow/tensorflow/issues/34429
-      ../cuda-10.2-no-bin2c-path.patch
     ];
 
     # On update, it can be useful to steal the changes from gentoo
@@ -147,7 +138,7 @@ let
       git
 
       # libs taken from system through the TF_SYS_LIBS mechanism
-      # grpc
+      grpc
       sqlite
       openssl
       jsoncpp
@@ -165,7 +156,6 @@ let
     ] ++ lib.optionals cudaSupport [
       cudatoolkit
       cudnn
-      nvidia_x11
     ] ++ lib.optionals mklSupport [
       mkl
     ] ++ lib.optionals stdenv.isDarwin [
@@ -182,24 +172,26 @@ let
     TF_SYSTEM_LIBS = lib.concatStringsSep "," [
       "absl_py"
       "astor_archive"
+      "astunparse_archive"
       "boringssl"
       # Not packaged in nixpkgs
       # "com_github_googleapis_googleapis"
       # "com_github_googlecloudplatform_google_cloud_cpp"
+      "com_github_grpc_grpc"
       "com_google_protobuf"
       "com_googlesource_code_re2"
       "curl"
       "cython"
       "double_conversion"
+      "enum34_archive"
       "flatbuffers"
+      "functools32_archive"
       "gast_archive"
-      # Lots of errors, requires an older version
-      # "grpc"
+      "gif"
       "hwloc"
       "icu"
-      "jpeg"
       "jsoncpp_git"
-      "keras_applications_archive"
+      "libjpeg_turbo"
       "lmdb"
       "nasm"
       # "nsync" # not packaged in nixpkgs
@@ -207,12 +199,14 @@ let
       "org_sqlite"
       "pasta"
       "pcre"
+      "png"
+      "pybind11"
       "six_archive"
       "snappy"
       "swig"
       "termcolor_archive"
       "wrapt"
-      "zlib_archive"
+      "zlib"
     ];
 
     INCLUDEDIR = "${includes_joined}/include";
@@ -235,13 +229,16 @@ let
     TF_CUDA_COMPUTE_CAPABILITIES = lib.concatStringsSep "," cudaCapabilities;
 
     postPatch = ''
-      # https://github.com/tensorflow/tensorflow/issues/20919
-      sed -i '/androidndk/d' tensorflow/lite/kernels/internal/BUILD
-
       # Tensorboard pulls in a bunch of dependencies, some of which may
       # include security vulnerabilities. So we make it optional.
       # https://github.com/tensorflow/tensorflow/issues/20280#issuecomment-400230560
       sed -i '/tensorboard >=/d' tensorflow/tools/pip_package/setup.py
+
+      # numpy 1.19 added in https://github.com/tensorflow/tensorflow/commit/75ea0b31477d6ba9e990e296bbbd8ca4e7eebadf.patch
+      sed -i 's/numpy >= 1.16.0, < 1.19.0/numpy >= 1.16.0/' tensorflow/tools/pip_package/setup.py
+
+      # bazel 3.3 should work just as well as bazel 3.1
+      rm -f .bazelversion
     '';
 
     preConfigure = let
@@ -272,15 +269,8 @@ let
       runHook postConfigure
     '';
 
-    # FIXME: Tensorflow uses dlopen() for CUDA libraries.
-    NIX_LDFLAGS = lib.optionalString cudaSupport "-lcudart -lcublas -lcufft -lcurand -lcusolver -lcusparse -lcudnn";
-
     hardeningDisable = [ "format" ];
 
-    bazelFlags = [
-      # temporary fixes to make the build work with bazel 0.27
-      "--incompatible_no_support_tools_in_action_inputs=false"
-    ];
     bazelBuildFlags = [
       "--config=opt" # optimize using the flags set in the configure phase
     ]
@@ -288,15 +278,17 @@ let
 
     bazelTarget = "//tensorflow/tools/pip_package:build_pip_package //tensorflow/tools/lib_package:libtensorflow";
 
+    removeRulesCC = false;
+
     fetchAttrs = {
       # So that checksums don't depend on these.
       TF_SYSTEM_LIBS = null;
 
       # cudaSupport causes fetch of ncclArchive, resulting in different hashes
       sha256 = if cudaSupport then
-        "1kqk1gx5g63kb2zdj392x6mnpbrmgqghrdv597aipn7s23xzj8pd"
+        "0pf8128chkm6fxnhd4956n6gvijlj00mjmvry33gq3xx3bayhs9g"
       else
-        "1plpcm2ydpajsrxdvmmpfy7l0gfdir78hap72w4k7ddm6d3rm2fv";
+        "0mkgss2nyk21zlj8hp24cs3dmpdnxk8qi6qq4hyc18lp82p09xwa";
     };
 
     buildAttrs = {
@@ -347,7 +339,7 @@ let
 
 in buildPythonPackage {
   inherit version pname;
-  disabled = isPy27 || (pythonAtLeast "3.8");
+  disabled = isPy27;
 
   src = bazel-build.python;
 
@@ -377,6 +369,8 @@ in buildPythonPackage {
     wrapt
     grpcio
     opt-einsum
+    astunparse
+    h5py
   ] ++ lib.optionals (!isPy3k) [
     mock
     future
@@ -392,6 +386,8 @@ in buildPythonPackage {
   postFixup = lib.optionalString cudaSupport ''
     find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
       addOpenGLRunpath "$lib"
+
+      patchelf --set-rpath "${cudatoolkit}/lib:${cudatoolkit.lib}/lib:${cudnn}/lib:${nccl}/lib:$(patchelf --print-rpath "$lib")" "$lib"
     done
   '';