nathan/nixpkgs
1
0
Fork 0
mirror of https://github.com/NixOS/nixpkgs.git synced 2026-08-26 02:05:02 +00:00
Code Issues Packages Projects Releases Wiki Activity

python3Packages.jax-cuda12-plugin: refactor src logic

Browse Source
This commit is contained in:
Gaetan Lepage
2026-08-22 12:38:49 +00:00
parent d928697767
commit 80232f258c
1 changed files with 32 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files

89
pkgs/development/python-modules/jax-cuda12-plugin/default.nix
View File

@@ -15,57 +15,29 @@ let
inherit (jaxlib) version;
inherit (jax-cuda12-pjrt) cudaLibPath;
getSrcFromPypi =
{
platform,
dist,
hash,
}:
fetchPypi {
inherit
version
platform
dist
hash
;
pname = "jax_cuda12_plugin";
format = "wheel";
python = dist;
abi = dist;
platforms = {
x86_64-linux = {
name = "manylinux_2_27_x86_64";
hashes = {
cp312 = "sha256-8BVPvyNrsiMab8pSBN0eQ6wMHPDeOgAg//kQPlP7hXI=";
cp313 = "sha256-2O/FH9SZiimQxGfmPDhs/ckYfH0QWCXFxKt0Pn1SLRE=";
cp314 = "sha256-a7jxBoIVF+NEiKUyDkX13vSkYo8qJUr6bix+7DVBm74=";
};
};
srcs = {
"3.12-x86_64-linux" = getSrcFromPypi {
platform = "manylinux_2_27_x86_64";
dist = "cp312";
hash = "sha256-8BVPvyNrsiMab8pSBN0eQ6wMHPDeOgAg//kQPlP7hXI=";
};
"3.12-aarch64-linux" = getSrcFromPypi {
platform = "manylinux_2_27_aarch64";
dist = "cp312";
hash = "sha256-8qOC5dv2YzB0C/TbOQPGEuhLJN4JkTx5U+0kA8KWtoY=";
};
"3.13-x86_64-linux" = getSrcFromPypi {
platform = "manylinux_2_27_x86_64";
dist = "cp313";
hash = "sha256-2O/FH9SZiimQxGfmPDhs/ckYfH0QWCXFxKt0Pn1SLRE=";
};
"3.13-aarch64-linux" = getSrcFromPypi {
platform = "manylinux_2_27_aarch64";
dist = "cp313";
hash = "sha256-6KjWb0PQCzXAQwNmsRfBthKDdk8NENOWWr/jn4coSG8=";
};
"3.14-x86_64-linux" = getSrcFromPypi {
platform = "manylinux_2_27_x86_64";
dist = "cp314";
hash = "sha256-a7jxBoIVF+NEiKUyDkX13vSkYo8qJUr6bix+7DVBm74=";
};
"3.14-aarch64-linux" = getSrcFromPypi {
platform = "manylinux_2_27_aarch64";
dist = "cp314";
hash = "sha256-PKEZswHZqTvLg3XGk3z3lJL/WVHKXLuLRvfjSvQtz9M=";
aarch64-linux = {
name = "manylinux_2_27_aarch64";
hashes = {
cp312 = "sha256-8qOC5dv2YzB0C/TbOQPGEuhLJN4JkTx5U+0kA8KWtoY=";
cp313 = "sha256-6KjWb0PQCzXAQwNmsRfBthKDdk8NENOWWr/jn4coSG8=";
cp314 = "sha256-PKEZswHZqTvLg3XGk3z3lJL/WVHKXLuLRvfjSvQtz9M=";
};
};
};
currentPlatform =
platforms.${stdenv.hostPlatform.system}
or (throw "jax-cuda12-plugin is not supported on ${stdenv.hostPlatform.system}");
dist = "cp${lib.replaceStrings [ "." ] [ "" ] python.pythonVersion}";
in
buildPythonPackage {
pname = "jax-cuda12-plugin";
@@ -73,10 +45,17 @@ buildPythonPackage {
pyproject = false;
__structuredAttrs = true;
src = (
srcs."${python.pythonVersion}-${stdenv.hostPlatform.system}"
or (throw "python${python.pythonVersion}Packages.jax-cuda12-plugin is not supported on ${stdenv.hostPlatform.system}")
);
src = fetchPypi {
pname = "jax_cuda12_plugin";
inherit version dist;
format = "wheel";
python = dist;
abi = dist;
platform = currentPlatform.name;
hash =
currentPlatform.hashes.${dist}
or (throw "python${python.pythonVersion}Packages.jax-cuda12-plugin is not supported");
};
nativeBuildInputs = [
autoPatchelfHook
@@ -119,11 +98,7 @@ buildPythonPackage {
sourceProvenance = [ lib.sourceTypes.binaryNativeCode ];
license = lib.licenses.asl20;
maintainers = with lib.maintainers; [ natsukium ];
# Keep in sync with `srcs` above
platforms = [
"x86_64-linux"
"aarch64-linux"
];
platforms = lib.attrNames platforms;
# see CUDA compatibility matrix
# https://jax.readthedocs.io/en/latest/installation.html#pip-installation-nvidia-gpu-cuda-installed-locally-harder
broken = !(lib.versionAtLeast cudaPackages.cudnn.version "9.1");