Change:
# Without --add-needed autoPatchelf forgets $ORIGIN on cuda>=8.0.5. postFixup = strings.optionalString (strings.versionAtLeast versionTriple "8.0.5") '' patchelf $out/lib/libcudnn.so --add-needed libcudnn_cnn_infer.so <ins> patchelf $out/lib/libcudnn_ops_infer.so --add-needed libcublas.so --add-needed libcublasLt.so </ins> ''; passthru = {

Discussion:
Jaxlib fails with the following message when Jax is trying to create new arrays. Steps to reproduce the behavior: the following to create an python environment with jax and jaxlib this environment jax and create an array: Jax should be able to create arrays on GPU. I tried with Cuda version 11 and 12, both resulted in the same error. And pytorch doesn't have this problem, so I think this is error is specific to jaxlib. <!-- Please @ people who are in the list of the offending package or module. If in doubt, check for whoever last touched something. --Please run and paste the result.
Experiencing the same problem!
I'm using and appending either or to before launching python seems to workaround this issue. edit: I'm using nixpkgs at the 23.05 tag.
I did some digging, and the problem seems to come from , the libraries are trying to load and but failed. They used to have and in DTNEEDED, but now they don't Here is an from 8.6.0 and 8.9.1 and just went missing, so autopatchelf won't help. I patched cudnn with And now jax/jaxlib can work without any further isssues. But I not sure if it is the correct way to fix it, might need some help from cudnn's maintainers.
I independently found the same solution and worked fine for doing some training in my end.
Can someone post a step-by-step fix for this?
Here is an pr that would fix this issue You could wait for this pr to get merged, or if in a hurry, you could use branch directly. This should be quite easy if using flake, just target the nixpkgs input to this specific branch. Example: now run will drop you into a shell with the fix cuda environment. Note that you will need to spend some time to compile cudnn as there is no cache available.
Thank you for posting that, unfortunately it didn't work for me: Also I'm not sure why it's using cuda-11 whereas on my nixos i have cuda 12
I think it's just a network issue, it failed to download cuda from nvidia.

And the nccl version is 2.16.5-1-cuda-11, so yes, it's cuda-11.
You're right. Still I fixed that problem but the original bug doesn't seem to be solved:
Well, it seems that pr doesn't fix all the problems. You can try my branch , I have tested it, and jax works.
my branch didn't work for you too? I tested mine again and here it works fine.
I'm able to reproduce. From ,
I get the same error also with 's branch.
I just tested it, no error on my side, I run some of my JAX programs and they all worked.
can you do inside the environment and then use jax. This will provide some debug info.
Yes:
I wonder if that's due to it being and now ?
I think and are just the same file, it's just a link . One thing I did notice is that your was on a different version. On my branch its , and on ogoid's branch, it's .
What are ya'll's nixpkgs config settings? Perhaps there's a difference in there?
How do i check that? This is my
This shows that you had in the first place, but somehow, when running python and jax, you were not loading the right version of jax (the one with patched cudnn). My guess is that the flake was evaluated in an impure mode and the in the system environment got leaked.
I didn't run with the flag. Is there another way I could have caused that?
Just putting this here for anyone in a similar boat (might only be relevant for future issues like this, since the fix for this exact issue seems merged). I got the CUDA-enabled (but not CUDA-bundled) pip version of JAX running in my current flake inside an FHS with this setup (the fix BillHuang2001 provided can work as an Override):