b12f/nixpkgs
1
0
Fork 0
Code Issues Pull requests Packages Projects Releases Wiki Activity

python3Packages.jaxlib: pin cudaPackages to 11.6

Browse source
This commit is contained in:
Samuel Ainsworth 2022-04-19 05:39:59 +00:00
parent 1344d5fe60
commit 92a001fa1c
1 changed files with 6 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
pkgs/top-level/python-packages.nix
View file

@ -4243,20 +4243,23 @@ in {
jaxlib-bin = callPackage ../development/python-modules/jaxlib/bin.nix {
cudaSupport = pkgs.config.cudaSupport or false;
inherit (self.tensorflow) cudaPackages;
# At the time of writing (2022-04-18), `cudaPackages.nccl` is broken, so we
# pin to `cudaPackages_11_6` instead.
cudaPackages = pkgs.cudaPackages_11_6;
};
jaxlib-build = callPackage ../development/python-modules/jaxlib {
# Some platforms don't have `cudaSupport` defined, hence the need for 'or false'.
cudaSupport = pkgs.config.cudaSupport or false;
inherit (self.tensorflow) cudaPackages;
# At the time of writing (2022-04-18), `cudaPackages.nccl` is broken, so we
# pin to `cudaPackages_11_6` instead.
cudaPackages = pkgs.cudaPackages_11_6;
};
jaxlib = self.jaxlib-build;
jaxlibWithCuda = self.jaxlib-build.override {
cudaSupport = true;
};
jaxlibWithoutCuda = self.jaxlib-build.override {