From 5bd4f7e74925527f8ddb5cc40e5fb44620d84090 Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Wed, 27 Nov 2024 16:47:44 +0100
Subject: [PATCH] python312Packages.tensorflow: rename generated wheel from
 tensorflow_cpu to tensorflow

---
 .../python-modules/tensorflow/bin.nix         | 148 ++++++++++--------
 1 file changed, 86 insertions(+), 62 deletions(-)

diff --git a/pkgs/development/python-modules/tensorflow/bin.nix b/pkgs/development/python-modules/tensorflow/bin.nix
index 7a5127582a64..cb709c98b720 100644
--- a/pkgs/development/python-modules/tensorflow/bin.nix
+++ b/pkgs/development/python-modules/tensorflow/bin.nix
@@ -54,7 +54,7 @@ let
   isCudaJetson = cudaSupport && cudaPackages.cudaFlags.isJetsonBuild;
   isCudaX64 = cudaSupport && stdenv.hostPlatform.isx86_64;
 in
-buildPythonPackage {
+buildPythonPackage rec {
   pname = "tensorflow" + lib.optionalString cudaSupport "-gpu";
   version = packages."${"version" + lib.optionalString isCudaJetson "_jetson"}";
   format = "wheel";
@@ -141,73 +141,97 @@ buildPythonPackage {
     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
-      ];
+    # When using the cpu-only wheel, the final package will be named `tensorflow_cpu`.
+    # Then, in each package requiring `tensorflow`, our pythonRuntimeDepsCheck will fail with:
+    # importlib.metadata.PackageNotFoundError: No package metadata was found for tensorflow
+    # Hence, we manually rename the package to `tensorflow`.
+    lib.optionalString ((builtins.match ".*tensorflow_cpu.*" src.url) != null) ''
+      (
+        cd $out/${python.sitePackages}
 
-      libpaths = [
-        (lib.getLib stdenv.cc.cc)
-        zlib
-      ];
+        dest="tensorflow-${version}.dist-info"
 
-      rpath = lib.makeLibraryPath (libpaths ++ cudapaths);
-    in
-    lib.optionalString stdenv.hostPlatform.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/"
-        "$out/${python.sitePackages}/tensorflow/core/kernels"
-        "$out/${python.sitePackages}/tensorflow/compiler/mlir/stablehlo/"
-        "$out/${python.sitePackages}/tensorflow/compiler/tf2tensorrt/"
-        "$out/${python.sitePackages}/tensorflow/compiler/tf2xla/ops/"
-        "$out/${python.sitePackages}/tensorflow/include/external/ml_dtypes/"
-        "$out/${python.sitePackages}/tensorflow/lite/experimental/microfrontend/python/ops/"
-        "$out/${python.sitePackages}/tensorflow/lite/python/analyzer_wrapper/"
-        "$out/${python.sitePackages}/tensorflow/lite/python/interpreter_wrapper/"
-        "$out/${python.sitePackages}/tensorflow/lite/python/metrics/"
-        "$out/${python.sitePackages}/tensorflow/lite/python/optimize/"
-        "$out/${python.sitePackages}/tensorflow/python/"
-        "$out/${python.sitePackages}/tensorflow/python/autograph/impl/testing"
-        "$out/${python.sitePackages}/tensorflow/python/client"
-        "$out/${python.sitePackages}/tensorflow/python/data/experimental/service"
-        "$out/${python.sitePackages}/tensorflow/python/framework"
-        "$out/${python.sitePackages}/tensorflow/python/grappler"
-        "$out/${python.sitePackages}/tensorflow/python/lib/core"
-        "$out/${python.sitePackages}/tensorflow/python/lib/io"
-        "$out/${python.sitePackages}/tensorflow/python/platform"
-        "$out/${python.sitePackages}/tensorflow/python/profiler/internal"
-        "$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}"
+        mv tensorflow_cpu-${version}.dist-info "$dest"
+
+        (
+          cd "$dest"
+
+          substituteInPlace METADATA \
+            --replace-fail "tensorflow_cpu" "tensorflow"
+          substituteInPlace RECORD \
+            --replace-fail "tensorflow_cpu" "tensorflow"
+        )
       )
+    ''
+    # 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.
+    + (
+      let
+        # rpaths we only need to add if CUDA is enabled.
+        cudapaths = lib.optionals cudaSupport [
+          cudatoolkit.out
+          cudatoolkit.lib
+          cudnn
+        ];
 
-      # 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"
+        libpaths = [
+          (lib.getLib stdenv.cc.cc)
+          zlib
+        ];
 
-      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 ''
-          addDriverRunpath "$lib"
-        ''}
-      done
-    '';
+        rpath = lib.makeLibraryPath (libpaths ++ cudapaths);
+      in
+      lib.optionalString stdenv.hostPlatform.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/"
+          "$out/${python.sitePackages}/tensorflow/core/kernels"
+          "$out/${python.sitePackages}/tensorflow/compiler/mlir/stablehlo/"
+          "$out/${python.sitePackages}/tensorflow/compiler/tf2tensorrt/"
+          "$out/${python.sitePackages}/tensorflow/compiler/tf2xla/ops/"
+          "$out/${python.sitePackages}/tensorflow/include/external/ml_dtypes/"
+          "$out/${python.sitePackages}/tensorflow/lite/experimental/microfrontend/python/ops/"
+          "$out/${python.sitePackages}/tensorflow/lite/python/analyzer_wrapper/"
+          "$out/${python.sitePackages}/tensorflow/lite/python/interpreter_wrapper/"
+          "$out/${python.sitePackages}/tensorflow/lite/python/metrics/"
+          "$out/${python.sitePackages}/tensorflow/lite/python/optimize/"
+          "$out/${python.sitePackages}/tensorflow/python/"
+          "$out/${python.sitePackages}/tensorflow/python/autograph/impl/testing"
+          "$out/${python.sitePackages}/tensorflow/python/client"
+          "$out/${python.sitePackages}/tensorflow/python/data/experimental/service"
+          "$out/${python.sitePackages}/tensorflow/python/framework"
+          "$out/${python.sitePackages}/tensorflow/python/grappler"
+          "$out/${python.sitePackages}/tensorflow/python/lib/core"
+          "$out/${python.sitePackages}/tensorflow/python/lib/io"
+          "$out/${python.sitePackages}/tensorflow/python/platform"
+          "$out/${python.sitePackages}/tensorflow/python/profiler/internal"
+          "$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}"
+        )
+
+        # 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 ''
+            addDriverRunpath "$lib"
+          ''}
+        done
+      ''
+    );
 
   # Upstream has a pip hack that results in bin/tensorboard being in both tensorflow
   # and the propagated input tensorboard, which causes environment collisions.