python3Packages.tensorflow_2: 2.3.2 -> 2.4.0

Also the following related changes:
* Removed Python 2 support because it's not supported by TF and related packages for a long time.
* Upgraded tensorboard and estimator packages to the required versions.
* Added extra plugins for tensorboard to support profiling.
* In the previous derivation versions, TF_SYSTEM_LIBS didn't have any effect because it was reset at repo fetching stage, so TF always used its own dependencies. Made TF_SYSTEM_LIBS actually work and fixed the errors caused by enabling it.
* Enabled tensorboard by default (but still keeping an option to disable it if needed).

pkgs/development/python-modules/flatbuffers/default.nix

@@ -8,6 +8,10 @@ buildPythonPackage rec {
 
   sourceRoot = "source/python";
 
+  # flatbuffers needs VERSION environment variable for setting the correct
+  # version, otherwise it uses the current date.
+  VERSION = "${version}";
+
   pythonImportsCheck = [ "flatbuffers" ];
 
   meta = flatbuffers.meta // {

pkgs/development/python-modules/gviz-api/default.nix

@@ -0,0 +1,26 @@
+{ lib, fetchPypi, buildPythonPackage
+, six
+}:
+
+buildPythonPackage rec {
+  pname = "gviz_api";
+  version = "1.9.0";
+  format = "wheel";
+
+  src = fetchPypi {
+    inherit pname version;
+    format = "wheel";
+    sha256 = "1yag559lpmwfdxpxn679a6ajifcbpgljr5n6k5b7rrj38k2xq7jg";
+  };
+
+  propagatedBuildInputs = [
+    six
+  ];
+
+  meta = with lib; {
+    description = "Python API for Google Visualization";
+    homepage = https://developers.google.com/chart/interactive/docs/dev/gviz_api_lib;
+    license = licenses.asl20;
+    maintainers = with maintainers; [ ndl ];
+  };
+}

pkgs/development/python-modules/tensorboard-plugin-profile/default.nix

@@ -0,0 +1,31 @@
+{ lib, fetchPypi, buildPythonPackage
+, gviz-api
+, protobuf
+, werkzeug
+}:
+
+buildPythonPackage rec {
+  pname = "tensorboard_plugin_profile";
+  version = "2.4.0";
+  format = "wheel";
+
+  src = fetchPypi {
+    inherit pname version;
+    format = "wheel";
+    python = "py3";
+    sha256 = "0z6dcjvkk3pzmmmjxi2ybawnfshz5qa3ga92kqj69ld1g9k3i9bj";
+  };
+
+  propagatedBuildInputs = [
+    gviz-api
+    protobuf
+    werkzeug
+  ];
+
+  meta = with lib; {
+    description = "Profile Tensorboard Plugin.";
+    homepage = http://tensorflow.org;
+    license = licenses.asl20;
+    maintainers = with maintainers; [ ndl ];
+  };
+}

pkgs/development/python-modules/tensorboard-plugin-wit/default.nix

@@ -0,0 +1,22 @@
+{ lib, fetchPypi, buildPythonPackage
+}:
+
+buildPythonPackage rec {
+  pname = "tensorboard_plugin_wit";
+  version = "1.7.0";
+  format = "wheel";
+
+  src = fetchPypi {
+    inherit pname version;
+    format = "wheel";
+    python = "py3";
+    sha256 = "0nv855qm2fav70lndsrv810pqgg41sbmd70fk86wk18ih825yxzf";
+  };
+
+  meta = with lib; {
+    description = "What-If Tool TensorBoard plugin.";
+    homepage = http://tensorflow.org;
+    license = licenses.asl20;
+    maintainers = with maintainers; [ ndl ];
+  };
+}

pkgs/development/python-modules/tensorflow-estimator/default.nix

@@ -6,13 +6,13 @@
 
 buildPythonPackage rec {
   pname = "tensorflow-estimator";
-  version = "2.3.0";
+  version = "2.4.0";
   format = "wheel";
 
   src = fetchPypi {
     pname = "tensorflow_estimator";
     inherit version format;
-    sha256 = "11n4sl9wfr00fv1i837b7a36ink86ggmlsgj7i06kcfc011h6pmp";
+    sha256 = "1w0pkcslm6934qqd6m5gxyjdlnb4pbl47k6s99wsh6dyvvr7nysv";
   };
 
   propagatedBuildInputs = [ mock numpy absl-py ];

pkgs/development/python-modules/tensorflow-tensorboard/default.nix

@@ -5,9 +5,10 @@
 , protobuf
 , grpcio
 , markdown
-, futures
 , absl-py
 , google-auth-oauthlib
+, tensorboard-plugin-wit
+, tensorboard-plugin-profile
 }:
 
 # tensorflow/tensorboard is built from a downloaded wheel, because
@@ -16,19 +17,16 @@
 
 buildPythonPackage rec {
   pname = "tensorflow-tensorboard";
-  version = "2.1.0";
+  version = "2.4.0";
   format = "wheel";
+  disabled = !isPy3k;
 
-  src = fetchPypi ({
+  src = fetchPypi {
     pname = "tensorboard";
     inherit version format;
-  } // (if isPy3k then {
     python = "py3";
-    sha256 = "1wpjdzhjpcdkyaahzd4bl71k4l30z5c55280ndiwj32hw70lxrp6";
+    sha256 = "0f17h6i398n8maam0r3rssqvdqnqbwjyf96nnhf482anm1iwdq6d";
-  } else {
+  };
-    python = "py2";
-    sha256 = "1f805839xa36wxb7xac9fyxzaww92vw4d50vs6g61wnlr4byp00w";
-  }));
 
   propagatedBuildInputs = [
     numpy
@@ -38,10 +36,12 @@ buildPythonPackage rec {
     grpcio
     absl-py
     google-auth-oauthlib
+    tensorboard-plugin-profile
+    tensorboard-plugin-wit
     # not declared in install_requires, but used at runtime
     # https://github.com/NixOS/nixpkgs/issues/73840
     wheel
-  ] ++ lib.optional (!isPy3k) futures;
+  ];
 
   # in the absence of a real test suite, run cli and imports
   checkPhase = ''

pkgs/development/python-modules/tensorflow/default.nix

@@ -1,24 +1,25 @@
-{ stdenv, pkgs, bazel_3, buildBazelPackage, lib, fetchFromGitHub, fetchpatch, symlinkJoin
+{ stdenv, bazel_3, buildBazelPackage, isPy3k, lib, fetchFromGitHub, symlinkJoin
 , addOpenGLRunpath
 # Python deps
-, buildPythonPackage, isPy3k, isPy27, pythonOlder, pythonAtLeast, python
+, buildPythonPackage, pythonOlder, pythonAtLeast, python
 # Python libraries
-, numpy, tensorflow-tensorboard_2, backports_weakref, mock, enum34, absl-py
+, numpy, tensorflow-tensorboard_2, absl-py
-, future, setuptools, wheel, keras-preprocessing, keras-applications, google-pasta
+, future, setuptools, wheel, keras-preprocessing, google-pasta
-, functools32
 , opt-einsum, astunparse, h5py
 , termcolor, grpcio, six, wrapt, protobuf, tensorflow-estimator_2
+, dill, flatbuffers-python, tblib, typing-extensions
 # Common deps
-, git, swig, which, binutils, glibcLocales, cython
+, git, pybind11, which, binutils, glibcLocales, cython, perl
 # Common libraries
-, jemalloc, openmpi, astor, gast, grpc, sqlite, openssl, jsoncpp, re2
+, jemalloc, openmpi, gast, grpc, sqlite, boringssl, jsoncpp
-, curl, snappy, flatbuffers, icu, double-conversion, libpng, libjpeg, giflib
+, curl, snappy, flatbuffers-core, lmdb-core, icu, double-conversion, libpng, libjpeg_turbo, giflib
 # Upsteam by default includes cuda support since tensorflow 1.15. We could do
 # 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, cudatoolkit ? null, cudnn ? null, nccl ? null
 , mklSupport ? false, mkl ? null
+, tensorboardSupport ? true
 # XLA without CUDA is broken
 , xlaSupport ? cudaSupport
 # Default from ./configure script
@@ -39,7 +40,7 @@ assert ! (stdenv.isDarwin && cudaSupport);
 assert mklSupport -> mkl != null;
 
 let
-  withTensorboard = pythonOlder "3.6";
+  withTensorboard = (pythonOlder "3.6") || tensorboardSupport;
 
   cudatoolkit_joined = symlinkJoin {
     name = "${cudatoolkit.name}-merged";
@@ -65,34 +66,40 @@ let
   includes_joined = symlinkJoin {
     name = "tensorflow-deps-merged";
     paths = [
-      pkgs.protobuf
       jsoncpp
     ];
   };
 
   tfFeature = x: if x then "1" else "0";
 
-  version = "2.3.2";
+  version = "2.4.0";
   variant = if cudaSupport then "-gpu" else "";
   pname = "tensorflow${variant}";
 
   pythonEnv = python.withPackages (_:
     [ # python deps needed during wheel build time (not runtime, see the buildPythonPackage part for that)
-      numpy
+      # This list can likely be shortened, but each trial takes multiple hours so won't bother for now.
-      keras-preprocessing
-      protobuf
-      wrapt
-      gast
-      astor
       absl-py
-      termcolor
+      astunparse
-      keras-applications
+      dill
+      flatbuffers-python
+      gast
+      google-pasta
+      grpcio
+      h5py
+      keras-preprocessing
+      numpy
+      opt-einsum
+      protobuf
       setuptools
+      six
+      tblib
+      tensorflow-estimator_2
+      tensorflow-tensorboard_2
+      termcolor
+      typing-extensions
       wheel
-  ] ++ lib.optionals (!isPy3k)
+      wrapt
-  [ future
-    functools32
-    mock
   ]);
 
   bazel-build = buildBazelPackage {
@@ -103,27 +110,21 @@ let
       owner = "tensorflow";
       repo = "tensorflow";
       rev = "v${version}";
-      sha256 = "sha256-ncwIkqLDqrB33pB9/FTlBklsIJUEvnDUmyAeUfufCFs=";
+      sha256 = "0yl06aypfxrcs35828xf04mkidz1x0j89v0q5h4d2xps1cb5rv3f";
     };
 
     patches = [
-      # Fixes for NixOS jsoncpp
+      # Relax too strict Python packages versions dependencies.
-      ./system-jsoncpp.patch
-
       ./relax-dependencies.patch
-
+      # Add missing `io_bazel_rules_docker` dependency.
-      # see https://github.com/tensorflow/tensorflow/issues/40688
+      ./workspace.patch
-      (fetchpatch {
-        url = "https://github.com/tensorflow/tensorflow/commit/75ea0b31477d6ba9e990e296bbbd8ca4e7eebadf.patch";
-        sha256 = "1xp1icacig0xm0nmb05sbrf4nw4xbln9fhc308birrv8286zx7wv";
-      })
     ];
 
     # On update, it can be useful to steal the changes from gentoo
     # https://gitweb.gentoo.org/repo/gentoo.git/tree/sci-libs/tensorflow
 
     nativeBuildInputs = [
-      swig which pythonEnv
+      which pythonEnv cython perl
     ] ++ lib.optional cudaSupport addOpenGLRunpath;
 
     buildInputs = [
@@ -135,19 +136,18 @@ let
       # libs taken from system through the TF_SYS_LIBS mechanism
       grpc
       sqlite
-      openssl
+      boringssl
       jsoncpp
-      pkgs.protobuf
       curl
+      pybind11
       snappy
-      flatbuffers
+      flatbuffers-core
       icu
       double-conversion
       libpng
-      libjpeg
+      libjpeg_turbo
       giflib
-      re2
+      lmdb-core
-      pkgs.lmdb
     ] ++ lib.optionals cudaSupport [
       cudatoolkit
       cudnn
@@ -173,10 +173,17 @@ let
       # "com_github_googleapis_googleapis"
       # "com_github_googlecloudplatform_google_cloud_cpp"
       "com_github_grpc_grpc"
-      "com_google_protobuf"
+      # Multiple issues with custom protobuf.
-      "com_googlesource_code_re2"
+      # First `com_github_googleapis` fails to configure. Can be worked around by disabling `com_github_googleapis`
+      # and related functionality, but then the next error is about "dangling symbolic link", and in general
+      # looks like that's only the beginning: see
+      # https://stackoverflow.com/questions/55578884/how-to-build-tensorflow-1-13-1-with-custom-protobuf
+      # "com_google_protobuf"
+      # Fails with the error: external/org_tensorflow/tensorflow/core/profiler/utils/tf_op_utils.cc:46:49: error: no matching function for call to 're2::RE2::FullMatch(absl::lts_2020_02_25::string_view&, re2::RE2&)'
+      # "com_googlesource_code_re2"
       "curl"
       "cython"
+      "dill_archive"
       "double_conversion"
       "enum34_archive"
       "flatbuffers"
@@ -198,8 +205,9 @@ let
       "pybind11"
       "six_archive"
       "snappy"
-      "swig"
+      "tblib_archive"
       "termcolor_archive"
+      "typing_extensions_archive"
       "wrapt"
       "zlib"
     ];
@@ -224,16 +232,13 @@ let
     TF_CUDA_COMPUTE_CAPABILITIES = lib.concatStringsSep "," cudaCapabilities;
 
     postPatch = ''
+      # bazel 3.3 should work just as well as bazel 3.1
+      rm -f .bazelversion
+    '' + lib.optionalString (!withTensorboard) ''
       # 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
+      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
     '';
 
     # https://github.com/tensorflow/tensorflow/pull/39470
@@ -277,16 +282,15 @@ let
     bazelTarget = "//tensorflow/tools/pip_package:build_pip_package //tensorflow/tools/lib_package:libtensorflow";
 
     removeRulesCC = false;
+    # Without this Bazel complaints about sandbox violations.
+    dontAddBazelOpts = true;
 
     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
-        "sha256-lEdPA9vhYO6vd5FgPMbFp2PkRvDBurPidYsxtJLXcbQ="
+        "0vyy1hv0jy5pqwvnc8pxb9isgnbw07c4a4d4wn61db00np114crz"
       else
-        "sha256-ZEY/bWo5M3Juw1x3CwhXYXZHD4q5LzWDlhgXnh4P95U=";
+        "0vczv5f9s4dxgwdkmf1y9b9ybh5d3y1nllqhb5q8aj9kq73izyn9";
     };
 
     buildAttrs = {
@@ -329,15 +333,13 @@ let
       license = licenses.asl20;
       maintainers = with maintainers; [ jyp abbradar ];
       platforms = with platforms; linux ++ darwin;
-      # The py2 build fails due to some issue importing protobuf. Possibly related to the fix in
+      broken = !(xlaSupport -> cudaSupport);
-      # https://github.com/akesandgren/easybuild-easyblocks/commit/1f2e517ddfd1b00a342c6abb55aef3fd93671a2b
-      broken = !(xlaSupport -> cudaSupport) || !isPy3k;
     };
   };
 
 in buildPythonPackage {
   inherit version pname;
-  disabled = isPy27;
+  disabled = !isPy3k;
 
   src = bazel-build.python;
 
@@ -354,27 +356,23 @@ in buildPythonPackage {
   # tensorflow/tools/pip_package/setup.py
   propagatedBuildInputs = [
     absl-py
-    astor
+    astunparse
+    dill
+    flatbuffers-python
     gast
     google-pasta
-    keras-applications
+    grpcio
+    h5py
     keras-preprocessing
     numpy
-    six
+    opt-einsum
     protobuf
+    six
+    tblib
     tensorflow-estimator_2
     termcolor
+    typing-extensions
     wrapt
-    grpcio
-    opt-einsum
-    astunparse
-    h5py
-  ] ++ lib.optionals (!isPy3k) [
-    mock
-    future
-    functools32
-  ] ++ lib.optionals (pythonOlder "3.4") [
-    backports_weakref enum34
   ] ++ lib.optionals withTensorboard [
     tensorflow-tensorboard_2
   ];

pkgs/development/python-modules/tensorflow/relax-dependencies.patch

@@ -1,16 +1,51 @@
 diff --git a/tensorflow/tools/pip_package/setup.py b/tensorflow/tools/pip_package/setup.py
-index 594e74f40c0..bfbf010144f 100644
+index 65133afdafe..8ef6364ff7e 100644
 --- a/tensorflow/tools/pip_package/setup.py
 +++ b/tensorflow/tools/pip_package/setup.py
-@@ -54,9 +54,9 @@ _VERSION = '2.3.1'
+@@ -75,23 +75,23 @@ if '--project_name' in sys.argv:
+ # comment the versioning scheme.
+ # NOTE: Please add test only packages to `TEST_PACKAGES` below.
  REQUIRED_PACKAGES = [
-     'absl-py >= 0.7.0',
+-    'absl-py ~= 0.10',
-     'astunparse == 1.6.3',
+-    'astunparse ~= 1.6.3',
+-    'flatbuffers ~= 1.12.0',
+-    'google_pasta ~= 0.2',
+-    'h5py ~= 2.10.0',
+-    'keras_preprocessing ~= 1.1.2',
+-    'numpy ~= 1.19.2',
+-    'opt_einsum ~= 3.3.0',
++    'absl-py >= 0.10',
++    'astunparse >= 1.6.3',
++    'flatbuffers >= 1.12.0',
++    'google_pasta >= 0.2',
++    'h5py >= 2.10.0',
++    'keras_preprocessing >= 1.1.2',
++    'numpy >= 1.19.1',
++    'opt_einsum >= 3.3.0',
+     'protobuf >= 3.9.2',
+-    'six ~= 1.15.0',
+-    'termcolor ~= 1.1.0',
+-    'typing_extensions ~= 3.7.4',
+-    'wheel ~= 0.35',
+-    'wrapt ~= 1.12.1',
++    'six >= 1.15.0',
++    'termcolor >= 1.1.0',
++    'typing_extensions >= 3.7.4',
++    'wheel >= 0.34.2',
++    'wrapt >= 1.12.1',
+     # These packages needs to be pinned exactly as newer versions are
+     # incompatible with the rest of the ecosystem
 -    'gast == 0.3.3',
 +    'gast >= 0.3.3',
-     'google_pasta >= 0.1.8',
+     # TensorFlow ecosystem packages that TF exposes API for
--    'h5py >= 2.10.0, < 2.11.0',
+     # These need to be in sync with the existing TF version
-+    'h5py >= 2.10.0',
+     # They are updated during the release process
-     'keras_preprocessing >= 1.1.1, < 1.2',
+@@ -118,7 +118,7 @@ if 'tf_nightly' in project_name:
-     # TODO(mihaimaruseac): numpy 1.19.0 has ABI breakage
+ # BoringSSL support.
-     # https://github.com/numpy/numpy/pull/15355
+ # See https://github.com/tensorflow/tensorflow/issues/17882.
+ if sys.byteorder == 'little':
+-  REQUIRED_PACKAGES.append('grpcio ~= 1.32.0')
++  REQUIRED_PACKAGES.append('grpcio >= 1.31.0')
+ 
+ 
+ # Packages which are only needed for testing code.

pkgs/development/python-modules/tensorflow/system-jsoncpp.patch

@@ -1,21 +0,0 @@
-diff --git a/third_party/systemlibs/jsoncpp.BUILD b/third_party/systemlibs/jsoncpp.BUILD
-index 526fd0c418..646f3fdcea 100644
---- a/third_party/systemlibs/jsoncpp.BUILD
-+++ b/third_party/systemlibs/jsoncpp.BUILD
-@@ -7,6 +7,7 @@ filegroup(
- 
- HEADERS = [
-     "include/json/autolink.h",
-+    "include/json/allocator.h",
-     "include/json/config.h",
-     "include/json/features.h",
-     "include/json/forwards.h",
-@@ -23,7 +24,7 @@ genrule(
-     cmd = """
-       for i in $(OUTS); do
-         i=$${i##*/}
--        ln -sf $(INCLUDEDIR)/jsoncpp/json/$$i $(@D)/include/json/$$i
-+        ln -sf $(INCLUDEDIR)/json/$$i $(@D)/include/json/$$i
-       done
-     """,
- )

pkgs/development/python-modules/tensorflow/workspace.patch

@@ -0,0 +1,18 @@
+diff --git a/WORKSPACE b/WORKSPACE
+index 9db1d9b80eb..c46f13f4ca4 100644
+--- a/WORKSPACE
++++ b/WORKSPACE
+@@ -12,6 +12,13 @@ http_archive(
+     ],
+ )
+ 
++http_archive(
++    name = "io_bazel_rules_docker",
++    sha256 = "1698624e878b0607052ae6131aa216d45ebb63871ec497f26c67455b34119c80",
++    strip_prefix = "rules_docker-0.15.0",
++    urls = ["https://github.com/bazelbuild/rules_docker/releases/download/v0.15.0/rules_docker-v0.15.0.tar.gz"],
++)
++
+ # Load tf_repositories() before loading dependencies for other repository so
+ # that dependencies like com_google_protobuf won't be overridden.
+ load("//tensorflow:workspace.bzl", "tf_repositories")

pkgs/top-level/python-packages.nix

@@ -2764,6 +2764,8 @@ in {
 
   guzzle_sphinx_theme = callPackage ../development/python-modules/guzzle_sphinx_theme { };
 
+  gviz-api = callPackage ../development/python-modules/gviz-api {};
+
   gwyddion = disabledIf isPy3k (toPythonModule (pkgs.gwyddion.override {
     pythonSupport = true;
     pythonPackages = self;
@@ -7372,6 +7374,10 @@ in {
 
   tenacity = callPackage ../development/python-modules/tenacity { };
 
+  tensorboard-plugin-profile = callPackage ../development/python-modules/tensorboard-plugin-profile { };
+
+  tensorboard-plugin-wit = callPackage ../development/python-modules/tensorboard-plugin-wit {};
+
   tensorboardx = callPackage ../development/python-modules/tensorboardx { };
 
   tensorflow-bin_2 = callPackage ../development/python-modules/tensorflow/bin.nix {
@@ -7388,9 +7394,10 @@ in {
     cudatoolkit = pkgs.cudatoolkit_11_0;
     cudnn = pkgs.cudnn_cudatoolkit_11_0;
     nccl = pkgs.nccl_cudatoolkit_11;
-    openssl = pkgs.openssl_1_1;
     inherit (pkgs.darwin.apple_sdk.frameworks) Foundation Security;
-    inherit (pkgs) flatbuffers;
+    flatbuffers-core = pkgs.flatbuffers;
+    flatbuffers-python = self.flatbuffers;
+    lmdb-core = pkgs.lmdb;
   };
 
   tensorflow-build = self.tensorflow-build_2;