diff --git a/pkgs/development/python-modules/tensorflow/bin.nix b/pkgs/development/python-modules/tensorflow/bin.nix index 79495d4e9ac2..f8476a2f1821 100644 --- a/pkgs/development/python-modules/tensorflow/bin.nix +++ b/pkgs/development/python-modules/tensorflow/bin.nix @@ -1,40 +1,44 @@ { - stdenv, lib, - fetchurl, + stdenv, buildPythonPackage, - isPy3k, - astor, + fetchurl, + + # buildInputs + llvmPackages, + + # build-system + distutils, + + # dependencies + ml-dtypes, + absl-py, + astunparse, + flatbuffers, gast, google-pasta, - wrapt, + grpcio, + h5py, + libclang, numpy, - six, - termcolor, + opt-einsum, packaging, protobuf, - absl-py, - grpcio, - mock, - scipy, - distutils, - wheel, - jax, - ml-dtypes, - opt-einsum, - tensorflow-estimator-bin, + requests, + six, tensorboard, + termcolor, + typing-extensions, + wrapt, + isPy3k, + mock, + config, cudaSupport ? config.cudaSupport, cudaPackages, zlib, python, addDriverRunpath, - astunparse, - flatbuffers, - h5py, - llvmPackages, - typing-extensions, }: # We keep this binary build for three reasons: @@ -50,7 +54,6 @@ let inherit (cudaPackages) cudatoolkit cudnn; isCudaJetson = cudaSupport && cudaPackages.cudaFlags.isJetsonBuild; - isCudaX64 = cudaSupport && stdenv.hostPlatform.isx86_64; in buildPythonPackage rec { pname = "tensorflow" + lib.optionalString cudaSupport "-gpu"; @@ -68,38 +71,36 @@ buildPythonPackage rec { buildInputs = [ llvmPackages.openmp ]; + build-system = [ + distutils + ]; + + nativeBuildInputs = + lib.optionals cudaSupport [ addDriverRunpath ] + ++ lib.optionals isCudaJetson [ cudaPackages.autoAddCudaCompatRunpath ]; + dependencies = [ + absl-py astunparse flatbuffers - typing-extensions - distutils + gast + google-pasta + grpcio + h5py + libclang + ml-dtypes + numpy + opt-einsum packaging protobuf - numpy - scipy - (if isCudaX64 then jax else ml-dtypes) - termcolor - grpcio + requests six - astor - absl-py - gast - opt-einsum - google-pasta - wrapt - tensorflow-estimator-bin tensorboard - h5py + termcolor + typing-extensions + wrapt ] ++ lib.optional (!isPy3k) mock; - build-system = - [ - distutils - wheel - ] - ++ lib.optionals cudaSupport [ addDriverRunpath ] - ++ lib.optionals isCudaJetson [ cudaPackages.autoAddCudaCompatRunpath ]; - preConfigure = '' unset SOURCE_DATE_EPOCH @@ -113,28 +114,6 @@ buildPythonPackage rec { mv "$f" "$(sed -E 's/(nv[0-9]+)\.0*([0-9]+)/\1.\2/' <<< "$f")" done - wheel unpack --dest unpacked ./*.whl - rm ./*.whl - ( - cd unpacked/tensorflow* - # Adjust dependency requirements: - # - Relax flatbuffers, gast, protobuf, tensorboard, and tensorflow-estimator version requirements that don't match what we have packaged - # - The purpose of python3Packages.libclang is not clear at the moment and we don't have it packaged yet - # - keras and tensorlow-io-gcs-filesystem will be considered as optional for now. - # - numpy was pinned to fix some internal tests: https://github.com/tensorflow/tensorflow/issues/60216 - sed -i *.dist-info/METADATA \ - -e "/Requires-Dist: flatbuffers/d" \ - -e "/Requires-Dist: gast/d" \ - -e "/Requires-Dist: keras/d" \ - -e "/Requires-Dist: libclang/d" \ - -e "/Requires-Dist: protobuf/d" \ - -e "/Requires-Dist: tensorboard/d" \ - -e "/Requires-Dist: tensorflow-estimator/d" \ - -e "/Requires-Dist: tensorflow-io-gcs-filesystem/d" \ - -e "s/Requires-Dist: numpy (.*)/Requires-Dist: numpy/" - ) - wheel pack ./unpacked/tensorflow* - popd '';