gador/nixpkgs
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

python312Packages.jaxlib-bin: 0.4.28 -> 0.4.38

Browse Source 
Diff: jax-ml/jax@jaxlib-v0.4.28...jaxlib-v0.4.38

Changelog: https://github.com/jax-ml/jax/releases/tag/jaxlib-v0.4.38
This commit is contained in:
natsukium
2024-12-29 20:36:35 +09:00
parent 3dbc7c385c
commit 214660c708
3 changed files with 67 additions and 147 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nixos/doc/manual/release-notes/rl-2505.section.md
+5
View File
@@ -258,6 +258,11 @@
- `siduck76-st` has been renamed to `st-snazzy`, like the project's [flake](https://github.com/siduck/st/blob/main/flake.nix).
- `python3Packages.jax` now directly depends on `python3Packages.jaxlib`.
As a result, packages that depend on jax no longer need to include jaxlib to their dependencies.
There is also a breaking change in the handling of CUDA. Instead of using a CUDA compatible jaxlib
as before, you can use plugins like `python3Packages.jax-cuda12-plugin`.
<!-- To avoid merge conflicts, consider adding your item at an arbitrary place in the list instead. -->
pkgs/development/python-modules/jaxlib/bin.nix
+61 -136
View File
@@ -6,46 +6,24 @@
{
absl-py,
autoAddDriverRunpath,
autoPatchelfHook,
buildPythonPackage,
config,
fetchPypi,
fetchurl,
flatbuffers,
jaxlib-build,
lib,
ml-dtypes,
python,
scipy,
stdenv,
# Options:
cudaSupport ? config.cudaSupport,
cudaPackages,
}:
let
inherit (cudaPackages) cudaVersion;
version = "0.4.28";
version = "0.4.38";
inherit (python) pythonVersion;
cudaLibPath = lib.makeLibraryPath (
with cudaPackages;
[
(lib.getLib cuda_cudart) # libcudart.so
(lib.getLib cuda_cupti) # libcupti.so
(lib.getLib cudnn) # libcudnn.so
(lib.getLib libcufft) # libcufft.so
(lib.getLib libcusolver) # libcusolver.so
(lib.getLib libcusparse) # libcusparse.so
]
);
# As of 2023-06-06, google/jax upstream is no longer publishing CPU-only wheels to their GCS bucket. Instead the
# official instructions recommend installing CPU-only versions via PyPI.
cpuSrcs =
srcs =
let
getSrcFromPypi =
{
@@ -68,182 +46,129 @@ let
};
in
{
"3.9-x86_64-linux" = getSrcFromPypi {
platform = "manylinux2014_x86_64";
dist = "cp39";
hash = "sha256-Slbr8FtKTBeRaZ2HTgcvP4CPCYa0AQsU+1SaackMqdw=";
};
"3.9-aarch64-darwin" = getSrcFromPypi {
platform = "macosx_11_0_arm64";
dist = "cp39";
hash = "sha256-sBVi7IrXVxm30DiXUkiel+trTctMjBE75JFjTVKCrTw=";
};
"3.9-x86_64-darwin" = getSrcFromPypi {
platform = "macosx_10_14_x86_64";
dist = "cp39";
hash = "sha256-T5jMg3srbG3P4Kt/+esQkxSSCUYRmqOvn6oTlxj/J4c=";
};
"3.10-x86_64-linux" = getSrcFromPypi {
platform = "manylinux2014_x86_64";
dist = "cp310";
hash = "sha256-47zcb45g+FVPQVwU2TATTmAuPKM8OOVGJ0/VRfh1dps=";
hash = "sha256-Ya7MuaJ8Z/24RQ9jVyQAGc1FEcudYqROR2R1bThIU60=";
};
"3.10-aarch64-linux" = getSrcFromPypi {
platform = "manylinux2014_aarch64";
dist = "cp310";
hash = "sha256-7hnBY6j98IOdTBi4il+/tOcxunxDdBbT5Ug+Vwu3ZOQ=";
};
"3.10-aarch64-darwin" = getSrcFromPypi {
platform = "macosx_11_0_arm64";
dist = "cp310";
hash = "sha256-8Djmi9ENGjVUcisLvjbmpEg4RDenWqnSg/aW8O2fjAk=";
hash = "sha256-MLL1LLUNdHNK8vR3wlM6elg+O7eyyKzes2Hud9lAV3o=";
};
"3.10-x86_64-darwin" = getSrcFromPypi {
platform = "macosx_10_14_x86_64";
dist = "cp310";
hash = "sha256-pCHSN/jCXShQFm0zRgPGc925tsJvUrxJZwS4eCKXvWY=";
hash = "sha256-VcGbnT8zpvxZ9kSqWiH7oCY5zN13bLSptVJmJfV4Of8=";
};
"3.11-x86_64-linux" = getSrcFromPypi {
platform = "manylinux2014_x86_64";
dist = "cp311";
hash = "sha256-Rc4PPIQM/4I2z/JsN/Jsn/B4aV+T4MFiwyDCgfUEEnU=";
hash = "sha256-J1H/cDfWqZfQvg53zEvjgcWp+buLMU7bdVwTpv2Wn0U=";
};
"3.11-aarch64-linux" = getSrcFromPypi {
platform = "manylinux2014_aarch64";
dist = "cp311";
hash = "sha256-Q9tYxMQnYnKWNmpWwQMY4fAPUDaQ4X+Uu0NEKT4ZleA=";
};
"3.11-aarch64-darwin" = getSrcFromPypi {
platform = "macosx_11_0_arm64";
dist = "cp311";
hash = "sha256-eThX+vN/Nxyv51L+pfyBH0NeQ7j7S1AgWERKf17M+Ck=";
hash = "sha256-P7DqrnNpFXr+y+rVCq8p5z/936d6IzXXIb2XlPPFEOQ=";
};
"3.11-x86_64-darwin" = getSrcFromPypi {
platform = "macosx_10_14_x86_64";
dist = "cp311";
hash = "sha256-L/gpDtx7ksfq5SUX9lSSYz4mey6QZ7rT5MMj0hPnfPU=";
hash = "sha256-tn/eq9bf7Qi3do873/tSEWAIX4MFZpvRl77vYdCN4Is=";
};
"3.12-x86_64-linux" = getSrcFromPypi {
platform = "manylinux2014_x86_64";
dist = "cp312";
hash = "sha256-RqGqhX9P7uikP8upXA4Kti1AwmzJcwtsaWVZCLo1n40=";
hash = "sha256-2tbAqWVnwG0IPARp/sQPIBIQsJk2W9aYvjGm0uyI/Vk=";
};
"3.12-aarch64-linux" = getSrcFromPypi {
platform = "manylinux2014_aarch64";
dist = "cp312";
hash = "sha256-SW9FsOABojQTCc0MdK8LZwU33O15wWjLIwz8x3PwqoY";
};
"3.12-aarch64-darwin" = getSrcFromPypi {
platform = "macosx_11_0_arm64";
dist = "cp312";
hash = "sha256-jdi//jhTcC9jzZJNoO4lc0pNGc1ckmvgM9dyun0cF10=";
hash = "sha256-8zvK/jLJelYuz2iU18QWdMgMCs3t+lQj1Jr1EUcUmHQ=";
};
"3.12-x86_64-darwin" = getSrcFromPypi {
platform = "macosx_10_14_x86_64";
dist = "cp312";
hash = "sha256-1sCaVFMpciRhrwVuc1FG0sjHTCKsdCaoRetp8ya096A=";
hash = "sha256-P+/qmF8EFYFvO7r9PwOkNwUCde+brJpywTFOFkSsV8E=";
};
"3.13-x86_64-linux" = getSrcFromPypi {
platform = "manylinux2014_x86_64";
dist = "cp313";
hash = "sha256-LOd7qM2pJZpLypevwcci5CkabEY6Y/jTcsbtyFEX1iU=";
};
"3.13-aarch64-linux" = getSrcFromPypi {
platform = "manylinux2014_aarch64";
dist = "cp313";
hash = "sha256-JIzKN3Hr8ksHD0lwE2TOraM+YTlEWwbHgsylrFrZK/Q=";
};
"3.13-aarch64-darwin" = getSrcFromPypi {
platform = "macosx_11_0_arm64";
dist = "cp313";
hash = "sha256-b+MmuK82Y4fdR8zzElg7Kxf+0ScSybdKZIsYoTy9ur8=";
};
"3.13-x86_64-darwin" = getSrcFromPypi {
platform = "macosx_10_14_x86_64";
dist = "cp313";
hash = "sha256-QeVa5YGKiC5XiehI9vFmh6wTK8+7Wl+hFKXRi3jQXy0=";
};
};
# Note that the prebuilt jaxlib binary requires specific version of CUDA to
# work. The cuda12 jaxlib binaries only works with CUDA 12.2, and cuda11
# jaxlib binaries only works with CUDA 11.8. This is why we need to find a
# binary that matches the provided cudaVersion.
gpuSrcVersionString = "cuda${cudaVersion}-${pythonVersion}";
# Find new releases at https://storage.googleapis.com/jax-releases
# When upgrading, you can get these hashes from prefetch.sh. See
# https://github.com/google/jax/issues/12879 as to why this specific URL is the correct index.
gpuSrcs = {
"cuda12.2-3.9" = fetchurl {
url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp39-cp39-manylinux2014_x86_64.whl";
hash = "sha256-d8LIl22gIvmWfoyKfXKElZJXicPQIZxdS4HumhwQGCw=";
};
"cuda12.2-3.10" = fetchurl {
url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp310-cp310-manylinux2014_x86_64.whl";
hash = "sha256-PXtWv+UEcMWF8LhWe6Z1UGkf14PG3dkJ0Iop0LiimnQ=";
};
"cuda12.2-3.11" = fetchurl {
url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp311-cp311-manylinux2014_x86_64.whl";
hash = "sha256-QO2WSOzmJ48VaCha596mELiOfPsAGLpGctmdzcCHE/o=";
};
"cuda12.2-3.12" = fetchurl {
url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp312-cp312-manylinux2014_x86_64.whl";
hash = "sha256-ixWMaIChy4Ammsn23/3cCoala0lFibuUxyUr3tjfFKU=";
};
};
in
buildPythonPackage {
pname = "jaxlib";
inherit version;
format = "wheel";
disabled =
!(
pythonVersion == "3.9"
|| pythonVersion == "3.10"
|| pythonVersion == "3.11"
|| pythonVersion == "3.12"
);
# See https://discourse.nixos.org/t/ofborg-does-not-respect-meta-platforms/27019/6.
src =
if !cudaSupport then
(cpuSrcs."${pythonVersion}-${stdenv.hostPlatform.system}"
or (throw "jaxlib-bin is not supported on ${stdenv.hostPlatform.system}")
)
else
gpuSrcs."${gpuSrcVersionString}" or (throw "jaxlib-bin: No gpuSrc for ${gpuSrcVersionString}");
src = (
srcs."${pythonVersion}-${stdenv.hostPlatform.system}"
or (throw "jaxlib-bin is not supported on ${stdenv.hostPlatform.system}")
);
# Prebuilt wheels are dynamically linked against things that nix can't find.
# Run `autoPatchelfHook` to automagically fix them.
nativeBuildInputs =
lib.optionals stdenv.hostPlatform.isLinux [ autoPatchelfHook ]
++ lib.optionals cudaSupport [ autoAddDriverRunpath ];
nativeBuildInputs = lib.optionals stdenv.hostPlatform.isLinux [ autoPatchelfHook ];
# Dynamic link dependencies
buildInputs = [ (lib.getLib stdenv.cc.cc) ];
# jaxlib contains shared libraries that open other shared libraries via dlopen
# and these implicit dependencies are not recognized by ldd or
# autoPatchelfHook. That means we need to sneak them into rpath. This step
# must be done after autoPatchelfHook and the automatic stripping of
# artifacts. autoPatchelfHook runs in postFixup and auto-stripping runs in the
# patchPhase.
preInstallCheck = lib.optional cudaSupport ''
shopt -s globstar
for file in $out/**/*.so; do
patchelf --add-rpath "${cudaLibPath}" "$file"
done
'';
propagatedBuildInputs = [
dependencies = [
absl-py
flatbuffers
ml-dtypes
scipy
];
# jaxlib looks for ptxas at runtime, eg when running `jax.random.PRNGKey(0)`.
# Linking into $out is the least bad solution. See
# * https://github.com/NixOS/nixpkgs/pull/164176#discussion_r828801621
# * https://github.com/NixOS/nixpkgs/pull/288829#discussion_r1493852211
# for more info.
postInstall = lib.optional cudaSupport ''
mkdir -p $out/${python.sitePackages}/jaxlib/cuda/bin
ln -s ${lib.getExe' cudaPackages.cuda_nvcc "ptxas"} $out/${python.sitePackages}/jaxlib/cuda/bin/ptxas
'';
pythonImportsCheck = [ "jaxlib" ];
inherit (jaxlib-build) pythonImportsCheck;
meta = with lib; {
meta = {
description = "Prebuilt jaxlib backend from PyPi";
homepage = "https://github.com/google/jax";
sourceProvenance = with sourceTypes; [ binaryNativeCode ];
license = licenses.asl20;
maintainers = with maintainers; [ samuela ];
platforms = [
"aarch64-darwin"
"x86_64-linux"
"x86_64-darwin"
];
broken =
!(cudaSupport -> lib.versionAtLeast cudaVersion "11.1")
|| !(cudaSupport -> lib.versionAtLeast cudaPackages.cudnn.version "8.2")
|| !(cudaSupport -> stdenv.hostPlatform.isLinux)
|| !(cudaSupport -> (gpuSrcs ? "cuda${cudaVersion}-${pythonVersion}"))
sourceProvenance = with lib.sourceTypes; [ binaryNativeCode ];
license = lib.licenses.asl20;
maintainers = with lib.maintainers; [ samuela ];
badPlatforms = [
# Fails at pythonImportsCheckPhase:
# ...-python-imports-check-hook.sh/nix-support/setup-hook: line 10: 28017 Illegal instruction: 4
# /nix/store/5qpssbvkzfh73xih07xgmpkj5r565975-python3-3.11.9/bin/python3.11 -c
# 'import os; import importlib; list(map(lambda mod: importlib.import_module(mod), os.environ["pythonImportsCheck"].split()))'
|| (stdenv.hostPlatform.isDarwin && stdenv.hostPlatform.isx86_64);
"x86_64-darwin"
];
};
}
pkgs/top-level/python-packages.nix
+1 -11
View File
@@ -6562,9 +6562,7 @@ self: super: with self; {
jax-jumpy = callPackage ../development/python-modules/jax-jumpy { };
jaxlib-bin = callPackage ../development/python-modules/jaxlib/bin.nix {
inherit (pkgs.config) cudaSupport;
};
jaxlib-bin = callPackage ../development/python-modules/jaxlib/bin.nix { };
jaxlib-build = callPackage ../development/python-modules/jaxlib rec {
# Some platforms don't have `cudaSupport` defined, hence the need for 'or false'.
@@ -6574,14 +6572,6 @@ self: super: with self; {
jaxlib = jaxlib-bin;
jaxlibWithCuda = self.jaxlib.override {
cudaSupport = true;
};
jaxlibWithoutCuda = self.jaxlib.override {
cudaSupport = false;
};
jaxopt = callPackage ../development/python-modules/jaxopt { };
jaxtyping = callPackage ../development/python-modules/jaxtyping { };