diff options
author | Samuel Ainsworth <skainsworth@gmail.com> | 2024-05-13 12:06:07 -0400 |
---|---|---|
committer | GitHub <noreply@github.com> | 2024-05-13 12:06:07 -0400 |
commit | 3a993d32444337caa4db6b85580d83825d6f596b (patch) | |
tree | dc2643f21beee8e13510ced65b7ef44c3c6e457c /pkgs/top-level | |
parent | 8fe1aa682515b81a58125b72afae246555d68bce (diff) | |
parent | 32d1bb1e3a15b047555a6f9c66d940e32f53d5fa (diff) |
Merge pull request #291705 from GaetanLepage/jax
python311Packages.{jax,jaxlib,jaxlib-bin}: 0.4.24 -> 0.4.28
Diffstat (limited to 'pkgs/top-level')
-rw-r--r-- | pkgs/top-level/all-packages.nix | 4 | ||||
-rw-r--r-- | pkgs/top-level/python-packages.nix | 5 |
2 files changed, 5 insertions, 4 deletions
diff --git a/pkgs/top-level/all-packages.nix b/pkgs/top-level/all-packages.nix index ebfa8ec827ac3..c4ea0ac511581 100644 --- a/pkgs/top-level/all-packages.nix +++ b/pkgs/top-level/all-packages.nix @@ -7125,10 +7125,6 @@ with pkgs; cudaPackages_12_3 = callPackage ./cuda-packages.nix { cudaVersion = "12.3"; }; cudaPackages_12 = cudaPackages_12_2; # Latest supported by cudnn - # Use the older cudaPackages for tensorflow and jax, as determined by cudnn - # compatibility: https://www.tensorflow.org/install/source#gpu - cudaPackagesGoogle = cudaPackages_11; - cudaPackages = recurseIntoAttrs cudaPackages_12; # TODO: move to alias diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index 924d3ad1c34f9..84d17a04983e0 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -14885,6 +14885,8 @@ self: super: with self; { tensorflow-bin = callPackage ../development/python-modules/tensorflow/bin.nix { inherit (pkgs.config) cudaSupport; + # https://www.tensorflow.org/install/source#gpu + cudaPackages = pkgs.cudaPackages_11; }; tensorflow-build = let @@ -14892,6 +14894,8 @@ self: super: with self; { protobufTF = pkgs.protobuf_21.override { abseil-cpp = pkgs.abseil-cpp_202301; }; + # https://www.tensorflow.org/install/source#gpu + cudaPackagesTF = pkgs.cudaPackages_11; grpcTF = (pkgs.grpc.overrideAttrs ( oldAttrs: rec { # nvcc fails on recent grpc versions, so we use the latest patch level @@ -14937,6 +14941,7 @@ self: super: with self; { inherit (pkgs.darwin.apple_sdk.frameworks) Foundation Security; flatbuffers-core = pkgs.flatbuffers; flatbuffers-python = self.flatbuffers; + cudaPackages = compat.cudaPackagesTF; protobuf-core = compat.protobufTF; protobuf-python = compat.protobuf-pythonTF; grpc = compat.grpcTF; |