diff options
Diffstat (limited to 'pkgs/development/python-modules/tensorflow/default.nix')
-rw-r--r-- | pkgs/development/python-modules/tensorflow/default.nix | 20 |
1 files changed, 11 insertions, 9 deletions
diff --git a/pkgs/development/python-modules/tensorflow/default.nix b/pkgs/development/python-modules/tensorflow/default.nix index c8e292e316744..d2f713082077b 100644 --- a/pkgs/development/python-modules/tensorflow/default.nix +++ b/pkgs/development/python-modules/tensorflow/default.nix @@ -19,8 +19,8 @@ # https://groups.google.com/a/tensorflow.org/forum/#!topic/developers/iRCt5m4qUz0 , config , cudaSupport ? config.cudaSupport -, cudaPackages ? { } -, cudaCapabilities ? cudaPackages.cudaFlags.cudaCapabilities +, cudaPackagesGoogle +, cudaCapabilities ? cudaPackagesGoogle.cudaFlags.cudaCapabilities , mklSupport ? false, mkl , tensorboardSupport ? true # XLA without CUDA is broken @@ -50,15 +50,15 @@ let # __ZN4llvm11SmallPtrSetIPKNS_10AllocaInstELj8EED1Ev in any of the # translation units, so the build fails at link time stdenv = - if cudaSupport then cudaPackages.backendStdenv + if cudaSupport then cudaPackagesGoogle.backendStdenv else if originalStdenv.isDarwin then llvmPackages_11.stdenv else originalStdenv; - inherit (cudaPackages) cudatoolkit nccl; + inherit (cudaPackagesGoogle) cudatoolkit nccl; # use compatible cuDNN (https://www.tensorflow.org/install/source#gpu) # cudaPackages.cudnn led to this: # https://github.com/tensorflow/tensorflow/issues/60398 cudnnAttribute = "cudnn_8_6"; - cudnn = cudaPackages.${cudnnAttribute}; + cudnn = cudaPackagesGoogle.${cudnnAttribute}; gentoo-patches = fetchzip { url = "https://dev.gentoo.org/~perfinion/patches/tensorflow-patches-2.12.0.tar.bz2"; hash = "sha256-SCRX/5/zML7LmKEPJkcM5Tebez9vv/gmE4xhT/jyqWs="; @@ -420,7 +420,9 @@ let x86_64-linux = if cudaSupport then "sha256-5VFMNHeLrUxW5RTr6EhT3pay9nWJ5JkZTGirDds5QkU=" else "sha256-KzgWV69Btr84FdwQ5JI2nQEsqiPg1/+TWdbw5bmxXOE="; - aarch64-linux = "sha256-9btXrNHqd720oXTPDhSmFidv5iaZRLjCVX8opmrMjXk="; + aarch64-linux = if cudaSupport + then "sha256-ty5+51BwHWE1xR4/0WcWTp608NzSAS/iiyN+9zx7/wI=" + else "sha256-9btXrNHqd720oXTPDhSmFidv5iaZRLjCVX8opmrMjXk="; x86_64-darwin = "sha256-gqb03kB0z2pZQ6m1fyRp1/Nbt8AVVHWpOJSeZNCLc4w="; aarch64-darwin = "sha256-WdgAaFZU+ePwWkVBhLzjlNT7ELfGHOTaMdafcAMD5yo="; }.${stdenv.hostPlatform.system} or (throw "unsupported system ${stdenv.hostPlatform.system}"); @@ -486,8 +488,8 @@ let broken = stdenv.isDarwin || !(xlaSupport -> cudaSupport) - || !(cudaSupport -> builtins.hasAttr cudnnAttribute cudaPackages) - || !(cudaSupport -> cudaPackages ? cudatoolkit); + || !(cudaSupport -> builtins.hasAttr cudnnAttribute cudaPackagesGoogle) + || !(cudaSupport -> cudaPackagesGoogle ? cudatoolkit); } // lib.optionalAttrs stdenv.isDarwin { timeout = 86400; # 24 hours maxSilent = 14400; # 4h, double the default of 7200s @@ -590,7 +592,7 @@ in buildPythonPackage { # Regression test for #77626 removed because not more `tensorflow.contrib`. passthru = { - inherit cudaPackages; + cudaPackages = cudaPackagesGoogle; deps = bazel-build.deps; libtensorflow = bazel-build.out; }; |