From d7fd7991864754566b92336e15fa205bf166bded Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Sat, 5 Sep 2026 16:50:01 +0200
Subject: [PATCH] python3Packages.jax-cuda12-plugin: cleanup postInstall

---
 .../python-modules/jax-cuda12-plugin/default.nix           | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/pkgs/development/python-modules/jax-cuda12-plugin/default.nix b/pkgs/development/python-modules/jax-cuda12-plugin/default.nix
index 02a949297aec..baa4ac8faafc 100644
--- a/pkgs/development/python-modules/jax-cuda12-plugin/default.nix
+++ b/pkgs/development/python-modules/jax-cuda12-plugin/default.nix
@@ -70,9 +70,10 @@ buildPythonPackage {
   # * https://github.com/NixOS/nixpkgs/pull/375186
   # for more info.
   postInstall = ''
-    mkdir -p $out/${python.sitePackages}/jax_cuda12_plugin/cuda/bin
-    ln -s ${lib.getExe' cudaPackages.cuda_nvcc "ptxas"} $out/${python.sitePackages}/jax_cuda12_plugin/cuda/bin
-    ln -s ${lib.getExe' cudaPackages.cuda_nvcc "nvlink"} $out/${python.sitePackages}/jax_cuda12_plugin/cuda/bin
+    export BINPATH="$out/${python.sitePackages}/jax_cuda12_plugin/cuda/bin"
+    mkdir -p $BINPATH
+    ln -s ${lib.getExe' cudaPackages.cuda_nvcc "ptxas"} $BINPATH/ptxas
+    ln -s ${lib.getExe' cudaPackages.cuda_nvcc "nvlink"} $BINPATH/nvlink
   '';
 
   # jax-cuda12-plugin contains shared libraries that open other shared libraries via dlopen