python.pkgs.tensorflow{,-bin,-estimator,-tensorboard}-2: Init at 2.1.0

Major breaking change from 1.x, so treating keeping both versions for now.

(adapted from 33f11be707)
(adapted from 9e8dea7986)

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

@@ -0,0 +1,68 @@
+{ lib, fetchPypi, buildPythonPackage, isPy3k
+, numpy
+, wheel
+, werkzeug
+, protobuf
+, grpcio
+, markdown
+, futures
+, absl-py
+, google-auth-oauthlib
+}:
+
+# tensorflow/tensorboard is built from a downloaded wheel, because
+# https://github.com/tensorflow/tensorboard/issues/719 blocks
+# buildBazelPackage.
+
+buildPythonPackage rec {
+  pname = "tensorflow-tensorboard";
+  version = "2.1.0";
+  format = "wheel";
+
+  src = fetchPypi ({
+    pname = "tensorboard";
+    inherit version;
+    format = "wheel";
+  } // (if isPy3k then {
+    python = "py3";
+    sha256 = "1wpjdzhjpcdkyaahzd4bl71k4l30z5c55280ndiwj32hw70lxrp6";
+  } else {
+    python = "py2";
+    sha256 = "1f805839xa36wxb7xac9fyxzaww92vw4d50vs6g61wnlr4byp00w";
+  }));
+
+  propagatedBuildInputs = [
+    numpy
+    werkzeug
+    protobuf
+    markdown
+    grpcio
+    absl-py
+    google-auth-oauthlib
+    # 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 = ''
+    $out/bin/tensorboard --help > /dev/null
+  '';
+
+  pythonImportsCheck = [
+    "tensorboard"
+    "tensorboard.backend"
+    "tensorboard.compat"
+    "tensorboard.data"
+    "tensorboard.plugins"
+    "tensorboard.summary"
+    "tensorboard.util"
+  ];
+
+  meta = with lib; {
+    description = "TensorFlow's Visualization Toolkit";
+    homepage = http://tensorflow.org;
+    license = licenses.asl20;
+    maintainers = with maintainers; [ abbradar ];
+  };
+}

pkgs/development/python-modules/tensorflow/1/bin.nix

@@ -15,7 +15,7 @@
 , grpcio
 , mock
 , backports_weakref
-, tensorflow-estimator
+, tensorflow-estimator_1
 , tensorflow-tensorboard
 , cudaSupport ? false
 , cudatoolkit ? null
@@ -72,7 +72,7 @@ in buildPythonPackage {
     gast
     google-pasta
     wrapt
-    tensorflow-estimator
+    tensorflow-estimator_1
     tensorflow-tensorboard
     keras-applications
     keras-preprocessing

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

@@ -3,11 +3,11 @@
 # Python deps
 , buildPythonPackage, isPy3k, isPy27, pythonOlder, pythonAtLeast, python
 # Python libraries
-, numpy, tensorflow-tensorboard, backports_weakref, mock, enum34, absl-py
+, numpy, tensorflow-tensorboard_1, backports_weakref, mock, enum34, absl-py
 , future, setuptools, wheel, keras-preprocessing, keras-applications, google-pasta
 , functools32
 , opt-einsum
-, termcolor, grpcio, six, wrapt, protobuf, tensorflow-estimator_1_15_1
+, termcolor, grpcio, six, wrapt, protobuf, tensorflow-estimator_1
 # Common deps
 , git, swig, which, binutils, glibcLocales, cython
 # Common libraries
@@ -105,9 +105,9 @@ let
 
     patches = [
       # Work around https://github.com/tensorflow/tensorflow/issues/24752
-      ./no-saved-proto.patch
+      ../no-saved-proto.patch
       # Fixes for NixOS jsoncpp
-      ./system-jsoncpp.patch
+      ../system-jsoncpp.patch
 
       # https://github.com/tensorflow/tensorflow/pull/29673
       (fetchpatch {
@@ -121,7 +121,6 @@ let
         sha256 = "1n9ypbrx36fc1kc9cz5b3p9qhg15xxhq4nz6ap3hwqba535nakfz";
       })
 
-
       (fetchpatch {
         # be compatible with gast >0.2 instead of only gast 0.2.2
         name = "gast-update.patch";
@@ -132,7 +131,7 @@ let
 
       # cuda 10.2 does not have "-bin2c-path" option anymore
       # https://github.com/tensorflow/tensorflow/issues/34429
-      ./cuda-10.2-no-bin2c-path.patch
+      ../cuda-10.2-no-bin2c-path.patch
     ];
 
     # On update, it can be useful to steal the changes from gentoo
@@ -373,7 +372,7 @@ in buildPythonPackage {
     numpy
     six
     protobuf
-    tensorflow-estimator_1_15_1
+    tensorflow-estimator_1
     termcolor
     wrapt
     grpcio
@@ -385,7 +384,7 @@ in buildPythonPackage {
   ] ++ lib.optionals (pythonOlder "3.4") [
     backports_weakref enum34
   ] ++ lib.optionals withTensorboard [
-    tensorflow-tensorboard
+    tensorflow-tensorboard_1
   ];
 
   nativeBuildInputs = lib.optional cudaSupport addOpenGLRunpath;

pkgs/development/python-modules/tensorflow/2/bin.nix

@@ -0,0 +1,179 @@
+{ stdenv
+, lib
+, fetchurl
+, buildPythonPackage
+, isPy3k, pythonOlder, isPy38
+, astor
+, gast
+, google-pasta
+, wrapt
+, numpy
+, six
+, termcolor
+, protobuf
+, absl-py
+, grpcio
+, mock
+, scipy
+, wheel
+, opt-einsum
+, backports_weakref
+, tensorflow-estimator_2
+, tensorflow-tensorboard
+, cudaSupport ? false
+, cudatoolkit ? null
+, cudnn ? null
+, nvidia_x11 ? null
+, zlib
+, python
+, symlinkJoin
+, keras-applications
+, keras-preprocessing
+, addOpenGLRunpath
+}:
+
+# We keep this binary build for two reasons:
+# - the source build doesn't work on Darwin.
+# - the source build is currently brittle and not easy to maintain
+
+assert cudaSupport -> cudatoolkit != null
+                   && cudnn != null
+                   && nvidia_x11 != null;
+
+# unsupported combination
+assert ! (stdenv.isDarwin && cudaSupport);
+
+let
+  packages = import ./binary-hashes.nix;
+
+  variant = if cudaSupport then "-gpu" else "";
+  pname = "tensorflow${variant}";
+
+in buildPythonPackage {
+  inherit pname;
+  inherit (packages) version;
+  format = "wheel";
+
+  disabled = isPy38;
+
+  src = let
+    pyVerNoDot = lib.strings.stringAsChars (x: if x == "." then "" else x) python.pythonVersion;
+    platform = if stdenv.isDarwin then "mac" else "linux";
+    unit = if cudaSupport then "gpu" else "cpu";
+    key = "${platform}_py_${pyVerNoDot}_${unit}";
+  in fetchurl packages.${key};
+
+  propagatedBuildInputs = [
+    protobuf
+    numpy
+    scipy
+    termcolor
+    grpcio
+    six
+    astor
+    absl-py
+    gast
+    opt-einsum
+    google-pasta
+    wrapt
+    tensorflow-estimator_2
+    tensorflow-tensorboard
+    keras-applications
+    keras-preprocessing
+  ] ++ lib.optional (!isPy3k) mock
+    ++ lib.optionals (pythonOlder "3.4") [ backports_weakref ];
+
+  nativeBuildInputs = [ wheel ] ++ lib.optional cudaSupport addOpenGLRunpath;
+
+  preConfigure = ''
+    unset SOURCE_DATE_EPOCH
+
+    # Make sure that dist and the wheel file are writable.
+    chmod u+rwx -R ./dist
+
+    pushd dist
+
+    # Unpack the wheel file.
+    wheel unpack --dest unpacked ./*.whl
+
+    # Tensorflow has a hard dependency on gast==0.2.2, but we relax it to
+    # gast==0.3.2.
+    substituteInPlace ./unpacked/tensorflow*/tensorflow_core/tools/pip_package/setup.py --replace "gast == 0.2.2" "gast == 0.3.2"
+    substituteInPlace ./unpacked/tensorflow*/tensorflow_*.dist-info/METADATA --replace "gast (==0.2.2)" "gast (==0.3.2)"
+
+    # Pack the wheel file back up.
+    wheel pack ./unpacked/tensorflow*
+
+    popd
+  '';
+
+  # Note that we need to run *after* the fixup phase because the
+  # libraries are loaded at runtime. If we run in preFixup then
+  # patchelf --shrink-rpath will remove the cuda libraries.
+  postFixup =
+    let
+      # rpaths we only need to add if CUDA is enabled.
+      cudapaths = lib.optionals cudaSupport [
+        cudatoolkit.out
+        cudatoolkit.lib
+        cudnn
+        nvidia_x11
+      ];
+
+      libpaths = [
+        stdenv.cc.cc.lib
+        zlib
+      ];
+
+      rpath = stdenv.lib.makeLibraryPath (libpaths ++ cudapaths);
+    in
+    lib.optionalString stdenv.isLinux ''
+      # This is an array containing all the directories in the tensorflow2
+      # package that contain .so files.
+      #
+      # TODO: Create this list programmatically, and remove paths that aren't
+      # actually needed.
+      rrPathArr=(
+        "$out/${python.sitePackages}/tensorflow_core/"
+        "$out/${python.sitePackages}/tensorflow_core/compiler/tf2tensorrt/"
+        "$out/${python.sitePackages}/tensorflow_core/compiler/tf2xla/ops/"
+        "$out/${python.sitePackages}/tensorflow_core/lite/experimental/microfrontend/python/ops/"
+        "$out/${python.sitePackages}/tensorflow_core/lite/python/interpreter_wrapper/"
+        "$out/${python.sitePackages}/tensorflow_core/lite/python/optimize/"
+        "$out/${python.sitePackages}/tensorflow_core/python/"
+        "$out/${python.sitePackages}/tensorflow_core/python/framework/"
+        "${rpath}"
+      )
+
+      # The the bash array into a colon-separated list of RPATHs.
+      rrPath=$(IFS=$':'; echo "''${rrPathArr[*]}")
+      echo "about to run patchelf with the following rpath: $rrPath"
+
+      find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
+        echo "about to patchelf $lib..."
+        chmod a+rx "$lib"
+        patchelf --set-rpath "$rrPath" "$lib"
+        ${lib.optionalString cudaSupport ''
+          addOpenGLRunpath "$lib"
+        ''}
+      done
+    '';
+
+  pythonImportsCheck = [
+    "tensorflow"
+    "tensorflow.keras"
+    "tensorflow.python"
+    "tensorflow.python.framework"
+  ];
+
+  meta = with stdenv.lib; {
+    description = "Computation using data flow graphs for scalable machine learning";
+    homepage = http://tensorflow.org;
+    license = licenses.asl20;
+    maintainers = with maintainers; [ jyp abbradar cdepillabout ];
+    platforms = [ "x86_64-linux" "x86_64-darwin" ];
+    # Python 2.7 build uses different string encoding.
+    # See https://github.com/NixOS/nixpkgs/pull/37044#issuecomment-373452253
+    broken = stdenv.isDarwin && !isPy3k;
+  };
+}

pkgs/development/python-modules/tensorflow/2/binary-hashes.nix

@@ -0,0 +1,51 @@
+{
+version = "2.1.0";
+linux_py_27_gpu = {
+  url = "https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-2.1.0-cp27-cp27mu-manylinux2010_x86_64.whl";
+  sha256 = "17lnhr7vdrls68c79n3sah5rpd0q1x2v5m84azvlyxxh2wpypfmb";
+};
+linux_py_27_cpu = {
+  url = "https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow_cpu-2.1.0-cp27-cp27mu-manylinux2010_x86_64.whl";
+  sha256 = "10lz3i4pcpgqrcbjmxm0n7k1gsqlpna3kdid902j2fy060cpi93z";
+};
+linux_py_35_gpu = {
+  url = "https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-2.1.0-cp35-cp35m-manylinux2010_x86_64.whl";
+  sha256 = "09s081n08dpmflwgir3zwzfijfpmahbh2gy5fn5bv5ll86g1szsy";
+};
+linux_py_35_cpu = {
+  url = "https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow_cpu-2.1.0-cp35-cp35m-manylinux2010_x86_64.whl";
+  sha256 = "1aa7v9fnvx03hqvhl3x3xcn41qy6qxw5xybg54ifjvvicp455c8l";
+};
+linux_py_36_gpu = {
+  url = "https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-2.1.0-cp36-cp36m-manylinux2010_x86_64.whl";
+  sha256 = "1dqp080ljbl9v3115vjp63ls0fimiwym6zxyanyhrlk8kwsq20zc";
+};
+linux_py_36_cpu = {
+  url = "https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow_cpu-2.1.0-cp36-cp36m-manylinux2010_x86_64.whl";
+  sha256 = "133z8anx7xm9rr5i9s9dwnp1wf06nr6s7q1lbs4lxpk6kn9nl480";
+};
+linux_py_37_gpu = {
+  url = "https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-2.1.0-cp37-cp37m-manylinux2010_x86_64.whl";
+  sha256 = "0yabl3xmcpr67w0zksqs3qc68nl9ax0vcd7w7b35nq8f65xl0ghy";
+};
+linux_py_37_cpu = {
+  url = "https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow_cpu-2.1.0-cp37-cp37m-manylinux2010_x86_64.whl";
+  sha256 = "04gngbngyg7p1gwx1q89my0cl8j7lq4kknqh51s2ynrix71zvsy6";
+};
+mac_py_27_cpu = {
+  url = "https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-2.1.0-cp27-cp27m-macosx_10_9_x86_64.whl";
+  sha256 = "1mprp72w5kk0lyjm2mh4lf57827xk3wsg28c4gizwm00ydfgacg6";
+};
+mac_py_35_cpu = {
+  url = "https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-2.1.0-cp35-cp35m-macosx_10_6_intel.whl";
+  sha256 = "1as7brf5ai6r7v1di9646jfrbnirpk2b0d1g29mn3shavb62kw8w";
+};
+mac_py_36_cpu = {
+  url = "https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-2.1.0-cp36-cp36m-macosx_10_9_x86_64.whl";
+  sha256 = "1v1rw9kjrskhcq1yas4ly2yfnzf2i1pjh6qg6zixfbkpkw7sw3wc";
+};
+mac_py_37_cpu = {
+  url = "https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-2.1.0-cp37-cp37m-macosx_10_9_x86_64.whl";
+  sha256 = "1hh4n0d97mrq35cmmsrnlmcv9vlswsyjy368lj3pda3y9dvck3rf";
+};
+}

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

@@ -0,0 +1,421 @@
+{ stdenv, pkgs, bazel_0_29, buildBazelPackage, lib, fetchFromGitHub, fetchpatch, symlinkJoin
+, addOpenGLRunpath
+# Python deps
+, buildPythonPackage, isPy3k, isPy27, pythonOlder, pythonAtLeast, python
+# Python libraries
+, numpy, tensorflow-tensorboard_2, backports_weakref, mock, enum34, absl-py
+, future, setuptools, wheel, keras-preprocessing, keras-applications, google-pasta
+, functools32
+, opt-einsum
+, termcolor, grpcio, six, wrapt, protobuf, tensorflow-estimator_2
+# Common deps
+, git, swig, which, binutils, glibcLocales, cython
+# Common libraries
+, jemalloc, openmpi, astor, gast, grpc, sqlite, openssl, jsoncpp, re2
+, curl, snappy, flatbuffers, icu, double-conversion, libpng, libjpeg, 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, nvidia_x11 ? null, cudatoolkit ? null, cudnn ? null, nccl ? null
+# XLA without CUDA is broken
+, xlaSupport ? cudaSupport
+# Default from ./configure script
+, cudaCapabilities ? [ "3.5" "5.2" ]
+, sse42Support ? builtins.elem (stdenv.hostPlatform.platform.gcc.arch or "default") ["westmere" "sandybridge" "ivybridge" "haswell" "broadwell" "skylake" "skylake-avx512"]
+, avx2Support  ? builtins.elem (stdenv.hostPlatform.platform.gcc.arch or "default") [                                     "haswell" "broadwell" "skylake" "skylake-avx512"]
+, fmaSupport   ? builtins.elem (stdenv.hostPlatform.platform.gcc.arch or "default") [                                     "haswell" "broadwell" "skylake" "skylake-avx512"]
+# Darwin deps
+, Foundation, Security
+}:
+
+assert cudaSupport -> nvidia_x11 != null
+                   && cudatoolkit != null
+                   && cudnn != null;
+
+# unsupported combination
+assert ! (stdenv.isDarwin && cudaSupport);
+
+let
+  withTensorboard = pythonOlder "3.6";
+
+  cudatoolkit_joined = symlinkJoin {
+    name = "${cudatoolkit.name}-merged";
+    paths = [
+      cudatoolkit.lib
+      cudatoolkit.out
+      # for some reason some of the required libs are in the targets/x86_64-linux
+      # directory; not sure why but this works around it
+      "${cudatoolkit}/targets/${stdenv.system}"
+    ];
+  };
+
+  cudatoolkit_cc_joined = symlinkJoin {
+    name = "${cudatoolkit.cc.name}-merged";
+    paths = [
+      cudatoolkit.cc
+      binutils.bintools # for ar, dwp, nm, objcopy, objdump, strip
+    ];
+  };
+
+  # Needed for _some_ system libraries, grep INCLUDEDIR.
+  includes_joined = symlinkJoin {
+    name = "tensorflow-deps-merged";
+    paths = [
+      pkgs.protobuf
+      jsoncpp
+    ];
+  };
+
+  tfFeature = x: if x then "1" else "0";
+
+  version = "2.1.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
+      keras-preprocessing
+      protobuf
+      wrapt
+      gast
+      astor
+      absl-py
+      termcolor
+      keras-applications
+      setuptools
+      wheel
+  ] ++ lib.optionals (!isPy3k)
+  [ future
+    functools32
+    mock
+  ]);
+
+  bazel-build = buildBazelPackage {
+    name = "${pname}-${version}";
+    bazel = bazel_0_29;
+
+    src = fetchFromGitHub {
+      owner = "tensorflow";
+      repo = "tensorflow";
+      rev = "v${version}";
+      sha256 = "1g79xi8yl4sjia8ysk9b7xfzrz83zy28v5dlb2wzmcf0k5pmz60p";
+    };
+
+    patches = [
+      # Work around https://github.com/tensorflow/tensorflow/issues/24752
+      ../no-saved-proto.patch
+      # Fixes for NixOS jsoncpp
+      ../system-jsoncpp.patch
+
+      (fetchpatch {
+        name = "backport-pr-18950.patch";
+        url = "https://github.com/tensorflow/tensorflow/commit/73640aaec2ab0234d9fff138e3c9833695570c0a.patch";
+        sha256 = "1n9ypbrx36fc1kc9cz5b3p9qhg15xxhq4nz6ap3hwqba535nakfz";
+      })
+
+      (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" ];
+      })
+
+      ./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
+    # https://gitweb.gentoo.org/repo/gentoo.git/tree/sci-libs/tensorflow
+
+    nativeBuildInputs = [
+      swig which pythonEnv
+    ] ++ lib.optional cudaSupport addOpenGLRunpath;
+
+    buildInputs = [
+      jemalloc
+      openmpi
+      glibcLocales
+      git
+
+      # libs taken from system through the TF_SYS_LIBS mechanism
+      # grpc
+      sqlite
+      openssl
+      jsoncpp
+      pkgs.protobuf
+      curl
+      snappy
+      flatbuffers
+      icu
+      double-conversion
+      libpng
+      libjpeg
+      giflib
+      re2
+      pkgs.lmdb
+    ] ++ lib.optionals cudaSupport [
+      cudatoolkit
+      cudnn
+      nvidia_x11
+    ] ++ lib.optionals stdenv.isDarwin [
+      Foundation
+      Security
+    ];
+
+    # arbitrarily set to the current latest bazel version, overly careful
+    TF_IGNORE_MAX_BAZEL_VERSION = true;
+
+    # Take as many libraries from the system as possible. Keep in sync with
+    # list of valid syslibs in
+    # https://github.com/tensorflow/tensorflow/blob/master/third_party/systemlibs/syslibs_configure.bzl
+    TF_SYSTEM_LIBS = lib.concatStringsSep "," [
+      "absl_py"
+      "astor_archive"
+      "boringssl"
+      # Not packaged in nixpkgs
+      # "com_github_googleapis_googleapis"
+      # "com_github_googlecloudplatform_google_cloud_cpp"
+      "com_google_protobuf"
+      "com_googlesource_code_re2"
+      "curl"
+      "cython"
+      "double_conversion"
+      "flatbuffers"
+      "gast_archive"
+      # Lots of errors, requires an older version
+      # "grpc"
+      "hwloc"
+      "icu"
+      "jpeg"
+      "jsoncpp_git"
+      "keras_applications_archive"
+      "lmdb"
+      "nasm"
+      # "nsync" # not packaged in nixpkgs
+      "opt_einsum_archive"
+      "org_sqlite"
+      "pasta"
+      "pcre"
+      "six_archive"
+      "snappy"
+      "swig"
+      "termcolor_archive"
+      "wrapt"
+      "zlib_archive"
+    ];
+
+    INCLUDEDIR = "${includes_joined}/include";
+
+    PYTHON_BIN_PATH = pythonEnv.interpreter;
+
+    TF_NEED_GCP = true;
+    TF_NEED_HDFS = true;
+    TF_ENABLE_XLA = tfFeature xlaSupport;
+
+    CC_OPT_FLAGS = " ";
+
+    # https://github.com/tensorflow/tensorflow/issues/14454
+    TF_NEED_MPI = tfFeature cudaSupport;
+
+    TF_NEED_CUDA = tfFeature cudaSupport;
+    TF_CUDA_PATHS = lib.optionalString cudaSupport "${cudatoolkit_joined},${cudnn},${nccl}";
+    GCC_HOST_COMPILER_PREFIX = lib.optionalString cudaSupport "${cudatoolkit_cc_joined}/bin";
+    GCC_HOST_COMPILER_PATH = lib.optionalString cudaSupport "${cudatoolkit_cc_joined}/bin/gcc";
+    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
+    '';
+
+    preConfigure = let
+      opt_flags = []
+        ++ lib.optionals sse42Support ["-msse4.2"]
+        ++ lib.optionals avx2Support ["-mavx2"]
+        ++ lib.optionals fmaSupport ["-mfma"];
+    in ''
+      patchShebangs configure
+
+      # dummy ldconfig
+      mkdir dummy-ldconfig
+      echo "#!${stdenv.shell}" > dummy-ldconfig/ldconfig
+      chmod +x dummy-ldconfig/ldconfig
+      export PATH="$PWD/dummy-ldconfig:$PATH"
+
+      export PYTHON_LIB_PATH="$NIX_BUILD_TOP/site-packages"
+      export CC_OPT_FLAGS="${lib.concatStringsSep " " opt_flags}"
+      mkdir -p "$PYTHON_LIB_PATH"
+
+      # To avoid mixing Python 2 and Python 3
+      unset PYTHONPATH
+    '';
+
+    configurePhase = ''
+      runHook preConfigure
+      ./configure
+      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
+    ];
+
+    bazelTarget = "//tensorflow/tools/pip_package:build_pip_package //tensorflow/tools/lib_package:libtensorflow";
+
+    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
+        "0hg3ysy644950a34j28hrb317cz8gcbb9n84d36wdailvnlshghb"
+      else
+        "1gy4pz9kn30wb9c4a9584fibb88c3h38y3dqa99yw1lbsbyyi28c";
+    };
+
+    buildAttrs = {
+      outputs = [ "out" "python" ];
+
+      preBuild = ''
+        patchShebangs .
+      '';
+
+      installPhase = ''
+        mkdir -p "$out"
+        tar -xf bazel-bin/tensorflow/tools/lib_package/libtensorflow.tar.gz -C "$out"
+        # Write pkgconfig file.
+        mkdir "$out/lib/pkgconfig"
+        cat > "$out/lib/pkgconfig/tensorflow.pc" << EOF
+        Name: TensorFlow
+        Version: ${version}
+        Description: Library for computation using data flow graphs for scalable machine learning
+        Requires:
+        Libs: -L$out/lib -ltensorflow
+        Cflags: -I$out/include/tensorflow
+        EOF
+
+        # build the source code, then copy it to $python (build_pip_package
+        # actually builds a symlink farm so we must dereference them).
+        bazel-bin/tensorflow/tools/pip_package/build_pip_package --src "$PWD/dist"
+        cp -Lr "$PWD/dist" "$python"
+      '';
+
+      postFixup = lib.optionalString cudaSupport ''
+        find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
+          addOpenGLRunpath "$lib"
+        done
+      '';
+    };
+
+    meta = with stdenv.lib; {
+      description = "Computation using data flow graphs for scalable machine learning";
+      homepage = http://tensorflow.org;
+      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
+      # https://github.com/akesandgren/easybuild-easyblocks/commit/1f2e517ddfd1b00a342c6abb55aef3fd93671a2b
+      broken = !(xlaSupport -> cudaSupport) || !isPy3k;
+    };
+  };
+
+in buildPythonPackage {
+  inherit version pname;
+  disabled = isPy27 || (pythonAtLeast "3.8");
+
+  src = bazel-build.python;
+
+  # Upstream has a pip hack that results in bin/tensorboard being in both tensorflow
+  # and the propagated input tensorflow-tensorboard, which causes environment collisions.
+  # Another possibility would be to have tensorboard only in the buildInputs
+  # https://github.com/tensorflow/tensorflow/blob/v1.7.1/tensorflow/tools/pip_package/setup.py#L79
+  postInstall = ''
+    rm $out/bin/tensorboard
+  '';
+
+  setupPyGlobalFlags = [ "--project_name ${pname}" ];
+
+  # tensorflow/tools/pip_package/setup.py
+  propagatedBuildInputs = [
+    absl-py
+    astor
+    gast
+    google-pasta
+    keras-applications
+    keras-preprocessing
+    numpy
+    six
+    protobuf
+    tensorflow-estimator_2
+    termcolor
+    wrapt
+    grpcio
+    opt-einsum
+  ] ++ lib.optionals (!isPy3k) [
+    mock
+    future
+    functools32
+  ] ++ lib.optionals (pythonOlder "3.4") [
+    backports_weakref enum34
+  ] ++ lib.optionals withTensorboard [
+    tensorflow-tensorboard_2
+  ];
+
+  nativeBuildInputs = lib.optional cudaSupport addOpenGLRunpath;
+
+  postFixup = lib.optionalString cudaSupport ''
+    find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
+      addOpenGLRunpath "$lib"
+    done
+  '';
+
+  # Actual tests are slow and impure.
+  # TODO try to run them anyway
+  # TODO better test (files in tensorflow/tools/ci_build/builds/*test)
+  checkPhase = ''
+    ${python.interpreter} <<EOF
+    # A simple "Hello world"
+    import tensorflow as tf
+    hello = tf.constant("Hello, world!")
+    tf.print(hello)
+
+    # Fit a simple model to random data
+    import numpy as np
+    np.random.seed(0)
+    tf.random.set_seed(0)
+    model = tf.keras.models.Sequential([
+        tf.keras.layers.Dense(1, activation="linear")
+    ])
+    model.compile(optimizer="sgd", loss="mse")
+
+    x = np.random.uniform(size=(1,1))
+    y = np.random.uniform(size=(1,))
+    model.fit(x, y, epochs=1)
+    EOF
+  '';
+  # Regression test for #77626 removed because not more `tensorflow.contrib`.
+
+  passthru.libtensorflow = bazel-build.out;
+
+  inherit (bazel-build) meta;
+}

pkgs/development/python-modules/tensorflow/2/lift-gast-restriction.patch

@@ -0,0 +1,11 @@
+diff --git a/tensorflow/tools/pip_package/setup.py b/tensorflow/tools/pip_package/setup.py
+index 992f2eae22..d9386f9b13 100644
+--- a/tensorflow/tools/pip_package/setup.py
++++ b/tensorflow/tools/pip_package/setup.py
+@@ -54,5 +54,5 @@ REQUIRED_PACKAGES = [
+     'enum34 >= 1.1.6;python_version<"3.4"',
+-    'gast == 0.2.2',
++    'gast >= 0.2.2',
+     'google_pasta >= 0.1.6',
+     'keras_applications >= 1.0.8',
+     'keras_preprocessing >= 1.0.5',

pkgs/development/python-modules/tensorflow/2/prefetcher.sh

@@ -0,0 +1,44 @@
+#!/usr/bin/env bash
+
+version=2.1.0
+
+# List of binary wheels for Tensorflow.  The most recent versions can be found
+# on the following page:
+# https://www.tensorflow.org/install/pip?lang=python3#package-location
+url_and_key_list=(
+  "linux_py_27_gpu https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-${version}-cp27-cp27mu-manylinux2010_x86_64.whl"
+  "linux_py_27_cpu https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow_cpu-${version}-cp27-cp27mu-manylinux2010_x86_64.whl"
+  "linux_py_35_gpu https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-${version}-cp35-cp35m-manylinux2010_x86_64.whl"
+  "linux_py_35_cpu https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow_cpu-${version}-cp35-cp35m-manylinux2010_x86_64.whl"
+  "linux_py_36_gpu https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-${version}-cp36-cp36m-manylinux2010_x86_64.whl"
+  "linux_py_36_cpu https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow_cpu-${version}-cp36-cp36m-manylinux2010_x86_64.whl"
+  "linux_py_37_gpu https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-${version}-cp37-cp37m-manylinux2010_x86_64.whl"
+  "linux_py_37_cpu https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow_cpu-${version}-cp37-cp37m-manylinux2010_x86_64.whl"
+  "mac_py_27_cpu https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-${version}-cp27-cp27m-macosx_10_9_x86_64.whl"
+  "mac_py_35_cpu https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-${version}-cp35-cp35m-macosx_10_6_intel.whl"
+  "mac_py_36_cpu https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-${version}-cp36-cp36m-macosx_10_9_x86_64.whl"
+  "mac_py_37_cpu https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-${version}-cp37-cp37m-macosx_10_9_x86_64.whl"
+)
+
+hashfile=binary-hashes.nix
+rm -f $hashfile
+echo "{" >> $hashfile
+echo "version = \"$version\";" >> $hashfile
+
+for url_and_key in "${url_and_key_list[@]}"; do
+  key=$(echo "$url_and_key" | cut -d' ' -f1)
+  url=$(echo "$url_and_key" | cut -d' ' -f2)
+
+  echo "prefetching ${url}..."
+  hash=$(nix-prefetch-url $url)
+
+  echo "$key = {" >> $hashfile
+  echo "  url = \"$url\";" >> $hashfile
+  echo "  sha256 = \"$hash\";" >> $hashfile
+  echo "};" >> $hashfile
+
+  echo
+done
+
+echo "}" >> $hashfile
+echo "done."

pkgs/top-level/python-packages.nix

@@ -6548,21 +6548,39 @@ in {
 
   zerobin = callPackage ../development/python-modules/zerobin { };
 
-  tensorflow-estimator = callPackage ../development/python-modules/tensorflow-estimator { };
-  tensorflow-estimator_1_15_1 = callPackage ../development/python-modules/tensorflow-estimator/1_15_1.nix { };
+  tensorflow-estimator = self.tensorflow-estimator_1;
+
+  tensorflow-estimator_1 = callPackage ../development/python-modules/tensorflow-estimator/1 { };
+
+  tensorflow-estimator_2 = callPackage ../development/python-modules/tensorflow-estimator/2 { };
 
   tensorflow-probability = callPackage ../development/python-modules/tensorflow-probability { };
 
-  tensorflow-tensorboard = callPackage ../development/python-modules/tensorflow-tensorboard { };
+  tensorflow-tensorboard = self.tensorflow-tensorboard_1;
 
-  tensorflow-bin = callPackage ../development/python-modules/tensorflow/bin.nix {
+  tensorflow-tensorboard_1 = callPackage ../development/python-modules/tensorflow-tensorboard/1 { };
+
+  tensorflow-tensorboard_2 = callPackage ../development/python-modules/tensorflow-tensorboard/2 { };
+
+  tensorflow-bin = self.tensorflow-bin_1;
+
+  tensorflow-bin_1 = callPackage ../development/python-modules/tensorflow/1/bin.nix {
     cudaSupport = pkgs.config.cudaSupport or false;
     inherit (pkgs.linuxPackages) nvidia_x11;
     cudatoolkit = pkgs.cudatoolkit_10;
     cudnn = pkgs.cudnn_cudatoolkit_10;
   };
 
-  tensorflow-build = callPackage ../development/python-modules/tensorflow {
+  tensorflow-bin_2 = callPackage ../development/python-modules/tensorflow/2/bin.nix {
+    cudaSupport = pkgs.config.cudaSupport or false;
+    inherit (pkgs.linuxPackages) nvidia_x11;
+    cudatoolkit = pkgs.cudatoolkit_10;
+    cudnn = pkgs.cudnn_cudatoolkit_10;
+  };
+
+  tensorflow-build = self.tensorflow-build_1;
+
+  tensorflow-build_1 = callPackage ../development/python-modules/tensorflow/1 {
     cudaSupport = pkgs.config.cudaSupport or false;
     inherit (pkgs.linuxPackages) nvidia_x11;
     cudatoolkit = pkgs.cudatoolkit_10;
@@ -6572,7 +6590,19 @@ in {
     inherit (pkgs.darwin.apple_sdk.frameworks) Foundation Security;
   };
 
-  tensorflow = self.tensorflow-build;
+  tensorflow-build_2 = callPackage ../development/python-modules/tensorflow/2 {
+    cudaSupport = pkgs.config.cudaSupport or false;
+    inherit (pkgs.linuxPackages) nvidia_x11;
+    cudatoolkit = pkgs.cudatoolkit_10;
+    cudnn = pkgs.cudnn_cudatoolkit_10;
+    nccl = pkgs.nccl_cudatoolkit_10;
+    openssl = pkgs.openssl_1_1;
+    inherit (pkgs.darwin.apple_sdk.frameworks) Foundation Security;
+  };
+
+  tensorflow = self.tensorflow_1;
+  tensorflow_1 = self.tensorflow-build_1;
+  tensorflow_2 = self.tensorflow-build_2;
 
   tensorflowWithoutCuda = self.tensorflow.override {
     cudaSupport = false;