Change:
arch = # KeyError: ('Linux', 'arm64') <del> if stdenv.hostPlatform.isLinux && stdenv.hostPlatform.linuxArch == "arm64" then "aarch64" else stdenv.hostPlatform.linuxArch; </del> <ins> if effectiveStdenv.hostPlatform.isLinux && effectiveStdenv.hostPlatform.linuxArch == "arm64" then "aarch64" else effectiveStdenv.hostPlatform.linuxArch; xla = effectiveStdenv.mkDerivation { pname = "xla-src"; version = "unstable"; src = fetchFromGitHub { owner = "openxla"; repo = "xla"; # Update this according to https://github.com/google/jax/blob/jaxlib-v${version}/third_party/xla/workspace.bzl. rev = "12eee889e1f2ad41e27d7b0e970cb92d282d3ec5"; hash = "sha256-68kjjgwYjRlcT0TVJo9BN6s+WTkdu5UMJqQcfHpBT90="; }; patches = [ # Resolves "could not convert ‘result’ from ‘SmallVector<[...],6>’ to # ‘SmallVector<[...],4>’" compilation error. See https://github.com/google/jax/issues/19814#issuecomment-1945141259. (fetchpatch { url = "https://github.com/openxla/xla/commit/7a614cd346594fc7ea2fe75570c9c53a4a444f60.patch"; hash = "sha256-RtuQTH8wzNiJcOtISLhf+gMlH1gg8hekvxEB+4wX6BM="; }) ]; dontBuild = true; # This is necessary for patchShebangs to know the right path to use. nativeBuildInputs = [ python ]; # Main culprits we're targeting are third_party/tsl/third_party/gpus/crosstool/clang/bin/*.tpl postPatch = '' patchShebangs . ''; installPhase = '' cp -r . $out ''; }; </ins> bazel-build = buildBazelPackage rec { name = "bazel-build-${pname}-${version}";

Discussion:
When creating a dev shell with + , the resulting environment cannot use CUDA with . Steps to reproduce the behavior: a Note that I am using the commit from nixos-23.11, which is the latest at this moment. However, I have tried the commit from a month ago and the problem was there as well. , run with the above flake. and Note that it reports is found, but does not seem to have CUDA enabled. The above can be found . It is expected that and run with CUDA and GPU appears in the device list. I am running this on my NixOS machine with a Nvidia GPU, where nvidia-driver is properly installed. I am also an intensive user and it works perfectly. Also I vaguely remember such combination works correctly a few months ago. I am not particularly sure about the maintainers so I might be wrong. Based on 's commit history and meta I am goingt to and Please run and paste the result.
Let me know if I can provide more useful information. Any suggestion on how I can continue debugging this would also be very helpful. Thanks in advance!
hmm interesting... ooc what gpu are you trying to use? have you confirmed that the cuda capabilities specified in your flake are compatible with your gpu?
Also please gist the outputs of demonstrating which implementation of jaxlib ends up being loaded, and whether a search for is attempted
On this machine it is a 3060 Ti. I have been using on this machine, with the samilar dev Shell, so I can confirm the cuda capabilities specified should work.
Thanks. Let me put the stderr log here: The log is very huge. I tried to search for and it did not come back with anything. I am trying to see how is loaded from it.
I think it is trying to load from , but I am not sure whether it is the GPU-enabled one or not. I tend to think the above is the GPU-enabled because if I specify in the dev shell, it loads instead.
Also just to confirm: what's the output of ? Sometimes multiple jaxlib versions can get pulled in as transitive dependencies (though they shouldn't if everything is packaged correctly). Then python can end up loading jaxlib-without-cuda instead of jaxlib-with-cuda depending on the ordering in PYTHONPATH.

From the looks of your flake this shouldn't be the case, but I'm guessing that this is a reproduction distilled from a more complex flake/project?
Yes this is distilled from a real world flake. It seems that in my dev shell is not set. However, running produces: I looked into the of the : Apparently there is only one being pulled in here.
Ok yeah that looks fine.. hmm
Btw, I tried two more things the capabilities constraint - still the same this on a 3090 machine, a 3080 machine and a 4090 machine. All of them produce exactly the same behavior.
Attempting to repro this but currently blocked on
Interesting. I had an attempt to run my nix-based ML project on an Ubuntu machine but wasn't able to get the nix dev shell to recognize the GPU a while ago. Have been using NixOS for development since then ...
Oh yeah it definitely works... mostly. This is the first real issue I've had with nixGL haha
Will give it a try! Meanwhile, I can try out any thoughts/ideas you might have to debug. Thanks!
Yeah my hope upon getting a repro is to bisect nixpkgs commits to find a known-good commit and then hopefully isolate the issue in a commit range. If you're feeling up to the task it would be quite useful! This would be a great use case for GPU-enabled tests and some CI infrastructure. tl;dr is that build has been green for a while but at some point it stopped actually using/recognizing GPUs.
Got it. This looks like something I can try. It would probably take a while as I'd imagine CUDA will be built again and again.
Let me try to build from a few commits on to start with.
Tried from about 3 months ago. Got different behavior: An older version has similar issue:
Also confirmed that can reproduce the probem: Similarly, from 2023-12-01 can reproduce the issue:
The above seems to suggest that, it is likely that in the past 3 months would not work with CUDA.
I found the a successful commit on : Please note that this is not a commit from and I believe it is using CUDA 11.7
Nice! Well I guess we have a (rough) commit window now to bisect on.

Btw, I'm curious: does work any better for you on master? We don't have a CUDA alias for that one, so you'll need to set
Let me try that. My understanding is that is about downloading the pre-built binary (wheel) and patch it, while is about spinning up to build the wheel by ourselves And for end users of , they should behave pretty much the same. Is my understanding correct? Thanks!
Tried . Seeing the following error:
Digged a little bit into this. On the successful branch 23.05 with CUDA 11.7: while on the failed branch 23.11, This seems to suggest that the we have in this case is not enabled for cuda support.
Yes, that's exactly right! Nice find! Huh, very weird. We're building with all the GPU/CUDA flags turned on AFAIK. Also, it's surprising to me that jaxlib-bin isn't working on your system since upstream builds that and we make only minor modifications
I checked their and found that there is a bazel flag called . Don't know what it is (maybe related to ?), but 's does not have it. I tried to add it and (more specifically, the inside it) cannot build any more. Sorry if I did not make it clear. can be successfully built on my system. The problem is similar to , as it cannot be used with CUDA. It just prompts a different error.
More on . Although it can successfully create the "client", will report not being able to find cuda drivers: This is on 23.11. The on has been upgraded to , which has the same problem of not being able to find CUDA:
Tracing the code and I realized that the problem is that this line , which in turn calls in failed for . Based on , I think the <versionhere should be . And it turns out that we already have for , which contains . Scratching my head again now ...
Yup, totally with you there. Now if jaxlib-bin didn't build successfully that'd be an easier fix... Debugging with jaxlib-bin can be convenient since it offers "known-good" binaries to work with. Hmm, ok this is where output ought to point the way.

But looking through the logs you posted earlier it, I'm not able to find the string anywhere... Would you mind posting the logs for the version? OOC how did you come to conclude that <versionis the failing lookup? It's a perfectly sensible conclusion since just about every CUDA package ever has to load it, but just curious re your methodology.
Ah sorry I should have provided more information here. So my journey starts with the error: and this quickly leads to the code in tensorflow: This means that the call to failed. This particular function can be found at , which is essentially . The function is , which is basically a call to , which in turn is just a call to of the here. Reading the above functions, we can see that the offending filename is in the format of . I could not figure out the because it is actually defined , which is in turn injected from . Because I had the impression that I was using , so I believe the version here is . This turned out to be wrong, see below. Thanks a lot for the suggestion! I tried the following Which produces log: Turns out that it is actually looking for . However, it is looking for ! Below are the selected part of the above log: I then went back to check for the : , and apparently it is downloading the prebuilt binary that is tagged with . Note that points to Cuda 11.8. I might be wrong, but is prebuilt binary tagged with supposed to work with CUDA 11 as well? The logic that finds and loads the files definitely is not compatible with CUDA 11.
For as long as tf is 2.13: We should add the comment. It's possible that jaxlib and tf actually come with different version constraints, I haven't checked. I guess one then has to pass it ( specifically) . Looks like jaxlib is trying to take over a part of the dynamic linker's job instead of just using DT_NEEDED and naturally this doesn't play out smoothly. They probably wanted to defer loading or something
I am trying adding an overlay like this: Seems that is now working. I am not sure whether this will have bad implications though. Yeah they might have some reason to defer it.
There is one problem though, if I have the following overlay: is ok, but will fail to build because of test failures.

See the log: Note that all the failures seem to be related to: I suspect this is jax tests that should run on CUDA but incorrectly ran on cpu - but I am not sure.
At this moment, with the help of and (Thanks!), I think I might have a solution to this. TLTR: We can now use (version 0.4.23) on , if you apply an overlay to point to . NOTE: usually you will also need alongside . It works on , but you will have to disable the unit tests and import test. See below in "what is still unclear". Detailed explanation: download a prebuilt binary, which defers loading of at runtime via dynamically. prebuilt binary is built against , meaning that will look for . , there is another check to make sure that the CUDA version is exactly (i.e. 12.1 and 12.0 will not work). , uses , which is designed for . The current needs instead of . However, if your environment does not require , just point to and everything should work out. with fixed above, I cannot build unless all the tests are disabled. It does not seem to affect the actual usage of because the test failure is specifically about not finding GPU-version of during test. This should be fixable but I am not sure how at this moment. we want to be able to use instead of to have more flexibility. The current build of seems to produce a library that does not have CUDA support. This remains a mystery and needs further investigation.
I propose that we template the "cudaXX" strings in based on the version of the inputted / instead of hardcoding. That will prevent this kind of issue in the future. It's worth noting that upstream releases builds for both cuda11 and cuda12, so we're in luck there. Unfortunately I'm of minimal use in terms of submitting PRs until is resolved. This reminds me that we should also migrate jaxlib-bin to use the redist packages at some point. Fixing the source build will be its own adventure. I suspect we're just missing some kind of flag that changed in upstream's build process.
Sounds good. Still, for the upstream releases, cuda11 = cuda 11.8, cuda12 = cuda 12.2.

The minor version is still important here and should be checked. Agree. It probably requires setting a development environment and build incrementally with to investigate. Blindly update flag and build the whole package again from scratch is quite inefficient ...
Sorry but I was not able to find the build for . I saw the ones under , but not under . Are they the same?
Search for "cuda11/jaxlib-0.4.23" and you'll find them. This link seems to be outdated, or at least not actively maintained by upstream. We should prob update the comment to point to instead.
I saw you're using that with Ubuntu? Did you try nixglhost?
oooh no i hadn't heard of nixglhost before... i'll give that a shot!
With that I was able to find the one for CUDA 11 (cudnn 8.6) and it worked out with this overlay: Need some advice on how to fix on . Based on suggestion, I am thinking about the version of cuda, assert on 11.8 / 12.2 will be set to if cuda is 11.8, or if cuda is 12.2 the file based on the above version Does that make sense?
Yup, exactly what I was thinking! Though I have a sneaking suspicion that CUDA versions outside of 11.8 / 12.2 might work. Sometimes they'll attempt to and then try , and so forth. So we should try seeing if other CUDA versions work too.
Thanks! I tried 12.0 and it won't work for prebuilt binary originally built against 12.2.
Ah ok maybe not then... damn EDIT: the JAX folks claim that this ought to be possible:
I still don't understand (and haven't got around to verify myself) whether is functional, and why do we keep focusing on instead?
is not functional, which is the original report of this issue: can be built However, the built (more specifically, ), was not able to work with CUDA. It seems that the problem is is actually not "withCuda". There might be some flags to be set during bazel build, but I was not able to find the reason. We then turn to as an alternative because we need a working at this moment.

Turns out that has a different issue, but is easier to fix.
I also read about that, but my experiment wasn't successful, for example, it will panic about JAX being built with but the current CUDA version is
Looking at for https://hercules- I see no mention of () Is it intentional we add to but not to the following? (again, I've no idea what I'm talking about)
Similar here, I am not that familiar with jax's bazel code. Would like to provide some information here. with CUDA 11.7 works, and at that time looks pretty much the same:
It seems that the version of cudnn is not a requirement. Binary built with 8.6 can work with 8.9 at runtime.
Well, there's definitely a bug in here somewhere. It is def possible that this is part of it. I'm guessing that upstream changed their build flags and we completed a version upgrade without realizing since we don't have a means to run GPU-enabled tests in CI/nixpkgs-review yet. is included in as a means to reduce the number of possible hashes that we have to record. The fetch is a superset of the CPU-only fetch, so we combine them for simplicity. is really handy in these situations where the source build is either broken or not on a version that the user wants. The source build is really quite unwieldy unfortunately. (Similar but not quite as bad as the tensorflow situation.)
I am working on the fix for . Not sure what is the best way to test it though. I am thinking about: the fix in a branch a that points to that branch with CUDA 11 and 12 on 3.9 through 3.12
Another approach, IMHO a little easier but totally up to you: Then you can run this for , , etc and and . That way you can run/test locally without committing changes.
The PR is created at
I am closing this as at least the version is fixed now. Thanks a lot for the help
Fixing the *-bin version is great, but let's keep this open until the source version is fixed as well
FWIW, the Guix version (there is only a source version) is here: It works fine with CUDA.

The definition of the CUDA-less version that it references is here: Maybe there's an obvious difference somewhere in the arguments.
Ok I think I've found the culprit. Somewhere along the way we removed from and/or . Oops! Working on a fix now... Currently blocked on AFAICT we do not package CUB and it is not provided in any of the cudaPackages that I've consulted thus far. If anyone has any ideas on the best way to proceed, do let me know!
Pretty sure it's part of , and we also used to distribute it as part of nvidia-thrust which we deleted
Ah, thanks for the tip ! I wasn't able to find or : Perhaps it has a different attribute name? Assuming a typo, I checked nccl which is already a dependency of jaxlib, but to no avail: Perhaps I'm missing something? Btw, draft PR started here:
Oop, found , but not finding that header file sadly:
Not ""?
Adding seemed to solve the problem, despite the difference in filenames. Not sure exactly what kind of wizardry they're doing to adjust header names, but it seems happy for now.
Looking at your issue upstream, I think the names match?
Oop, yeah you're right. I was confusing that with something else
This issue has been mentioned on NixOS Discourse. There might be relevant details there:
If you're using NixOS is there a platform wide config to always have the cuda libraries enabled?
I'm trying to compile jax using the commit , which appears to have closed this issue. When I tried to run jax on CUDA using this commit, I would need to compile jaxlib from scratch. During the compilation, I get a hash mismatch error in a library that is compiling. I even tried to use a later commit to get the following output: This is my And this is my (for the newer commit):
This seems like a support question, you're more likely to find help on NixOS Discourse or Matrix. But the answer is and this will trigger a lot of rebuilds Interesting, this might be a reproducibility issue with the way we fetch bazel dependencies, in which case this deserves a separate issue