JAX now detects CUDA.

Switching entirely to nixpkgs/unstable and using a proper overlay to prevent JAX from compiling jaxlib was what made it work.
novaia · Mar 21, 2024 · dab3992 · dab3992
1 parent 5bd71be
commit dab3992
Showing 1 changed file with 6 additions and 14 deletions.
diff --git a/flake.nix b/flake.nix
@@ -7,18 +7,9 @@
         flake-utils.url = "github:numtide/flake-utils";
     };
     outputs = inputs@{ self, nixpkgs, nixpkgs-unstable, flake-utils, ... }:
-    flake-utils.lib.eachSystem [ "x86_64-linux" ] (system: let
-        inherit (nixpkgs) lib;
-        unstableCudaPkgs = import nixpkgs-unstable {
-            inherit system;
-            config = {
-                allowUnfree = true;
-                cudaSupport = true;
-            };
-        };
-    in {
+    flake-utils.lib.eachSystem [ "x86_64-linux" ] (system: {
         devShells = let
-            pyVer = "310";
+            pyVer = "311";
             py = "python${pyVer}";
             overlays = [
                 (final: prev: {
@@ -35,24 +26,25 @@
                                 nativeCheckInputs = [];
                                 pythonImportsCheck = [];
                                 pytestFlagsArray = [];
+                                passthru.tests = [];
                                 doCheck = false;
                             });
                         };
                     };
                 })
             ];
-            stableJaxPkgs = import nixpkgs {
+            unstableCudaPkgs = import nixpkgs-unstable {
                 inherit system overlays;
                 config = {
                     allowUnfree = true;
                     cudaSupport = true;
                 };
             };
         in rec {
-            default = stableJaxPkgs.mkShell {
+            default = unstableCudaPkgs.mkShell {
                 name = "cuda";
                 buildInputs = [
-                    (stableJaxPkgs.${py}.withPackages (pyp: with pyp; [
+                    (unstableCudaPkgs.${py}.withPackages (pyp: with pyp; [
                         jax
                         jaxlib-bin
                     ]))