diff --git a/pkgs/development/python-modules/jaxlib/bin.nix b/pkgs/development/python-modules/jaxlib/bin.nix index 7e6b00429df..0929831e32a 100644 --- a/pkgs/development/python-modules/jaxlib/bin.nix +++ b/pkgs/development/python-modules/jaxlib/bin.nix @@ -120,9 +120,15 @@ buildPythonPackage rec { done ''; - # pip dependencies and optionally cudatoolkit. Note that cudatoolkit is - # necessary since jaxlib looks for "ptxas" in $PATH. - propagatedBuildInputs = [ absl-py flatbuffers scipy ] ++ lib.optional cudaSupport cudatoolkit_11; + propagatedBuildInputs = [ absl-py flatbuffers scipy ]; + + # Note that cudatoolkit is snecessary since jaxlib looks for "ptxas" in $PATH. + # See https://github.com/NixOS/nixpkgs/pull/164176#discussion_r828801621 for + # more info. + postInstall = lib.optional cudaSupport '' + mkdir -p $out/bin + ln -s ${cudatoolkit_11}/bin/ptxas $out/bin/ptxas + ''; pythonImportsCheck = [ "jaxlib" ]; diff --git a/pkgs/development/python-modules/jaxlib/default.nix b/pkgs/development/python-modules/jaxlib/default.nix index 664e109719a..363bfe56134 100644 --- a/pkgs/development/python-modules/jaxlib/default.nix +++ b/pkgs/development/python-modules/jaxlib/default.nix @@ -259,7 +259,13 @@ buildPythonPackage { src = "${bazel-build}/jaxlib-${version}-cp${builtins.replaceStrings ["."] [""] python.pythonVersion}-none-manylinux2010_${stdenv.targetPlatform.linuxArch}.whl"; + # Note that cudatoolkit is necessary since jaxlib looks for "ptxas" in $PATH. + # See https://github.com/NixOS/nixpkgs/pull/164176#discussion_r828801621 for + # more info. postInstall = lib.optionalString cudaSupport '' + mkdir -p $out/bin + ln -s ${cudatoolkit}/bin/ptxas $out/bin/ptxas + find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do addOpenGLRunpath "$lib" patchelf --set-rpath "${cudatoolkit}/lib:${cudatoolkit.lib}/lib:${cudnn}/lib:${nccl}/lib:$(patchelf --print-rpath "$lib")" "$lib"