From e5ef2ddb74d51ec31419a8df52eefeda0a30f106 Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Wed, 15 May 2024 07:56:56 +0200
Subject: [PATCH] python311Packages.tensorflow-bin: 2.15.0 -> 2.16.1

Changelog: https://github.com/tensorflow/tensorflow/releases/tag/v2.16.1
---
 .../python-modules/tensorflow/bin.nix         | 16 ++++--
 .../tensorflow/binary-hashes.nix              | 50 ++++++++++++-------
 .../python-modules/tensorflow/prefetcher.sh   | 25 ++++++----
 3 files changed, 57 insertions(+), 34 deletions(-)

diff --git a/pkgs/development/python-modules/tensorflow/bin.nix b/pkgs/development/python-modules/tensorflow/bin.nix
index 65ca45f05666..feed6046d707 100644
--- a/pkgs/development/python-modules/tensorflow/bin.nix
+++ b/pkgs/development/python-modules/tensorflow/bin.nix
@@ -31,6 +31,7 @@
 , astunparse
 , flatbuffers
 , h5py
+, llvmPackages
 , typing-extensions
 }:
 
@@ -56,7 +57,11 @@ in buildPythonPackage {
     key = "${platform}_py_${pyVerNoDot}_${unit}";
   in fetchurl (packages.${key} or {});
 
-  propagatedBuildInputs = [
+  buildInputs = [
+    llvmPackages.openmp
+  ];
+
+  dependencies = [
     astunparse
     flatbuffers
     typing-extensions
@@ -81,7 +86,7 @@ in buildPythonPackage {
     h5py
   ] ++ lib.optional (!isPy3k) mock;
 
-  nativeBuildInputs = [ wheel ] ++ lib.optionals cudaSupport [ addOpenGLRunpath ];
+  build-system = [ wheel ] ++ lib.optionals cudaSupport [ addOpenGLRunpath ];
 
   preConfigure = ''
     unset SOURCE_DATE_EPOCH
@@ -91,7 +96,6 @@ in buildPythonPackage {
 
     pushd dist
 
-    orig_name="$(echo ./*.whl)"
     wheel unpack --dest unpacked ./*.whl
     rm ./*.whl
     (
@@ -113,7 +117,6 @@ in buildPythonPackage {
         -e "s/Requires-Dist: numpy (.*)/Requires-Dist: numpy/"
     )
     wheel pack ./unpacked/tensorflow*
-    mv *.whl $orig_name # avoid changes to the _os_arch.whl suffix
 
     popd
   '';
@@ -168,6 +171,7 @@ in buildPythonPackage {
         "$out/${python.sitePackages}/tensorflow/python/saved_model"
         "$out/${python.sitePackages}/tensorflow/python/util"
         "$out/${python.sitePackages}/tensorflow/tsl/python/lib/core"
+        "$out/${python.sitePackages}/tensorflow.libs/"
         "${rpath}"
       )
 
@@ -206,5 +210,9 @@ in buildPythonPackage {
     license = licenses.asl20;
     maintainers = with maintainers; [ jyp abbradar ];
     platforms = [ "x86_64-linux" "x86_64-darwin" ];
+    # Cannot import tensortfow on python 3.12 as it still dependends on distutils:
+    # ModuleNotFoundError: No module named 'distutils'
+    # https://github.com/tensorflow/tensorflow/issues/58073
+    broken = pythonAtLeast "3.12";
   };
 }
diff --git a/pkgs/development/python-modules/tensorflow/binary-hashes.nix b/pkgs/development/python-modules/tensorflow/binary-hashes.nix
index 9204d1652d03..754a84b8ef51 100644
--- a/pkgs/development/python-modules/tensorflow/binary-hashes.nix
+++ b/pkgs/development/python-modules/tensorflow/binary-hashes.nix
@@ -1,39 +1,51 @@
 {
-version = "2.15.0";
+version = "2.16.1";
 linux_py_39_cpu = {
-  url = "https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow_cpu-2.15.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl";
-  sha256 = "01m5wn9wpywd2cscklr2h8wdsf3wzg11dhgl0nfj7a997v2bvlyg";
+  url = "https://storage.googleapis.com/tensorflow/versions/2.16.1/tensorflow_cpu-2.16.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl";
+  sha256 = "0j4311fav3b4wdhw9wb3zvc3grvgxddl28i4kwnbsvvjb191pcs5";
 };
 linux_py_310_cpu = {
-  url = "https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow_cpu-2.15.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl";
-  sha256 = "0qjahvq9gd89cx2dwkb6gk3s947czmasp9cx1pqyw1x5f113s2zn";
+  url = "https://storage.googleapis.com/tensorflow/versions/2.16.1/tensorflow_cpu-2.16.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl";
+  sha256 = "0b8q3r2ql7rknlshiv3bh3hj9zll0dn6w0wdxnqgg18cljys5nhh";
 };
 linux_py_311_cpu = {
-  url = "https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow_cpu-2.15.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl";
-  sha256 = "0w3jz7rpbf4r8hm91rjh666qyqflqy6x619fas1iqcgnfx8pjv7l";
+  url = "https://storage.googleapis.com/tensorflow/versions/2.16.1/tensorflow_cpu-2.16.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl";
+  sha256 = "12vv2jxybdfa8nq7r2y53vl7lvmaacb1xmy2368hkwqm2ld4mplg";
+};
+linux_py_312_cpu = {
+  url = "https://storage.googleapis.com/tensorflow/versions/2.16.1/tensorflow_cpu-2.16.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl";
+  sha256 = "0vvv3cr10r6kgkalk3hydqz18y4lp990pm4vwq8mbwdi101m59mf";
 };
 linux_py_39_gpu = {
-  url = "https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow-2.15.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl";
-  sha256 = "0ckij6jffsi6yvjfiysnwrfblhqm5g1f8dkw4i7n3dwndnanba1z";
+  url = "https://storage.googleapis.com/tensorflow/versions/2.16.1/tensorflow-2.16.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl";
+  sha256 = "17v85433ns360rcccx7pnkhjzd003cr8p6rj51wqb3gl7gkldf83";
 };
 linux_py_310_gpu = {
-  url = "https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow-2.15.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl";
-  sha256 = "1y0cxsps4rkidc2c3apaf7vd9vsnjg5lpcq30bnrnzcvvx1inq7d";
+  url = "https://storage.googleapis.com/tensorflow/versions/2.16.1/tensorflow-2.16.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl";
+  sha256 = "0w3jd7siamvhgd7si59zpv3rdpq3s2ii3hhcdaimji66axv095hy";
 };
 linux_py_311_gpu = {
-  url = "https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow-2.15.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl";
-  sha256 = "1z7bd4yrdxa7rq4k1fybbwnyhxf4rm9nwx15bpk2m53c5hmyrs6y";
+  url = "https://storage.googleapis.com/tensorflow/versions/2.16.1/tensorflow-2.16.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl";
+  sha256 = "014liy5wnbi2d2lnr8lqhd7221a4a1v6gzih7nv5qfnf1h86234k";
+};
+linux_py_312_gpu = {
+  url = "https://storage.googleapis.com/tensorflow/versions/2.16.1/tensorflow-2.16.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl";
+  sha256 = "0qd09bxfnad969y2h3mlzib51y6xcsqipq1w278k1fl8ayxky4ja";
 };
 mac_py_39_cpu = {
-  url = "https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-2.15.0-cp39-cp39-macosx_10_15_x86_64.whl";
-  sha256 = "1kqb5h4vs4jlhlqjmqvcg6jivccwqj1rgklrgrhvw09wzzhxvz1c";
+  url = "https://storage.googleapis.com/tensorflow/versions/2.16.1/tensorflow-2.16.1-cp39-cp39-macosx_10_15_x86_64.whl";
+  sha256 = "0dkcd1v0ivjp6v5pngg0x461wci94nxvwprc6jwgi7rqgyx3zkz9";
 };
 mac_py_310_cpu = {
-  url = "https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-2.15.0-cp310-cp310-macosx_10_15_x86_64.whl";
-  sha256 = "0wmawdfhva9ixal7468c0rrx8qd7nzgkxy6i9k2s7cqn8c7qw94v";
+  url = "https://storage.googleapis.com/tensorflow/versions/2.16.1/tensorflow-2.16.1-cp310-cp310-macosx_10_15_x86_64.whl";
+  sha256 = "0rnlxrb4q2cgqaifalhpji15g1f3m9flqfvfj2p2sv3lvmbg2ydb";
 };
 mac_py_311_cpu = {
-  url = "https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-2.15.0-cp311-cp311-macosx_10_15_x86_64.whl";
-  sha256 = "0cigk43fabps7vskz4pfgmca3wa0545b15qridyqdbyp5ri1c1qy";
+  url = "https://storage.googleapis.com/tensorflow/versions/2.16.1/tensorflow-2.16.1-cp311-cp311-macosx_10_15_x86_64.whl";
+  sha256 = "1xc9lis61adk2y22a82i3rvc0j7qznf4fqwg2pd1665zlhybi9gq";
+};
+mac_py_312_cpu = {
+  url = "https://storage.googleapis.com/tensorflow/versions/2.16.1/tensorflow-2.16.1-cp312-cp312-macosx_10_15_x86_64.whl";
+  sha256 = "10g8y02ncp5l7f133vank6h01palc445h6s9jndmmy7vm33c7jh9";
 };
 }
diff --git a/pkgs/development/python-modules/tensorflow/prefetcher.sh b/pkgs/development/python-modules/tensorflow/prefetcher.sh
index 136d982f2f60..a7c54f22fb57 100755
--- a/pkgs/development/python-modules/tensorflow/prefetcher.sh
+++ b/pkgs/development/python-modules/tensorflow/prefetcher.sh
@@ -1,22 +1,25 @@
 #!/usr/bin/env bash
 
-version="2.15.0"
+version="2.16.1"
 
-bucket="https://storage.googleapis.com/tensorflow"
+bucket="https://storage.googleapis.com/tensorflow/versions/${version}"
 
 # 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_39_cpu $bucket/linux/cpu/tensorflow_cpu-${version}-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl"
-"linux_py_310_cpu $bucket/linux/cpu/tensorflow_cpu-${version}-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl"
-"linux_py_311_cpu $bucket/linux/cpu/tensorflow_cpu-${version}-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl"
-"linux_py_39_gpu $bucket/linux/gpu/tensorflow-${version}-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl"
-"linux_py_310_gpu $bucket/linux/gpu/tensorflow-${version}-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl"
-"linux_py_311_gpu $bucket/linux/gpu/tensorflow-${version}-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl"
-"mac_py_39_cpu $bucket/mac/cpu/tensorflow-${version}-cp39-cp39-macosx_10_15_x86_64.whl"
-"mac_py_310_cpu $bucket/mac/cpu/tensorflow-${version}-cp310-cp310-macosx_10_15_x86_64.whl"
-"mac_py_311_cpu $bucket/mac/cpu/tensorflow-${version}-cp311-cp311-macosx_10_15_x86_64.whl"
+"linux_py_39_cpu $bucket/tensorflow_cpu-${version}-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl"
+"linux_py_310_cpu $bucket/tensorflow_cpu-${version}-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl"
+"linux_py_311_cpu $bucket/tensorflow_cpu-${version}-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl"
+"linux_py_312_cpu $bucket/tensorflow_cpu-${version}-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl"
+"linux_py_39_gpu $bucket/tensorflow-${version}-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl"
+"linux_py_310_gpu $bucket/tensorflow-${version}-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl"
+"linux_py_311_gpu $bucket/tensorflow-${version}-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl"
+"linux_py_312_gpu $bucket/tensorflow-${version}-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl"
+"mac_py_39_cpu $bucket/tensorflow-${version}-cp39-cp39-macosx_10_15_x86_64.whl"
+"mac_py_310_cpu $bucket/tensorflow-${version}-cp310-cp310-macosx_10_15_x86_64.whl"
+"mac_py_311_cpu $bucket/tensorflow-${version}-cp311-cp311-macosx_10_15_x86_64.whl"
+"mac_py_312_cpu $bucket/tensorflow-${version}-cp312-cp312-macosx_10_15_x86_64.whl"
 )
 
 hashfile=binary-hashes.nix