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

python3Packages.jax-cuda12-pjrt: cleanup postInstall

Browse Source
This commit is contained in:
Gaetan Lepage
2026-09-05 16:49:56 +02:00
parent 8b188f6dc2
commit e85f7ece4e
1 changed files with 6 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
pkgs/development/python-modules/jax-cuda12-pjrt/default.nix
View File

@@ -71,10 +71,12 @@ buildPythonPackage (finalAttrs: {
# * https://github.com/NixOS/nixpkgs/pull/288829#discussion_r1493852211
# for more info.
postInstall = ''
mkdir -p $out/${python.sitePackages}/jax_plugins/nvidia/cuda_nvcc/bin
ln -s ${lib.getExe' cudaPackages.cuda_nvcc "ptxas"} $out/${python.sitePackages}/jax_plugins/nvidia/cuda_nvcc/bin/ptxas
ln -s ${lib.getExe' cudaPackages.cuda_nvcc "nvlink"} $out/${python.sitePackages}/jax_plugins/nvidia/cuda_nvcc/bin/nvlink
ln -s ${cudaPackages.cuda_nvcc}/nvvm $out/${python.sitePackages}/jax_plugins/nvidia/cuda_nvcc/nvvm
export OUTPATH="$out/${python.sitePackages}/jax_plugins/nvidia/cuda_nvcc"
export BINPATH="$OUTPATH/bin"
mkdir -p $BINPATH
ln -s ${lib.getExe' cudaPackages.cuda_nvcc "ptxas"} $BINPATH/ptxas
ln -s ${lib.getExe' cudaPackages.cuda_nvcc "nvlink"} $BINPATH/nvlink
ln -s ${cudaPackages.cuda_nvcc}/nvvm $OUTPATH/nvvm
'';
# jax-cuda12-pjrt contains shared libraries that open other shared libraries via dlopen