In reply to @glepage:matrix.org
Where is the right spot to set cudaSupport = true for a flake based nixos install ?

 35       perSystem = {                                                                                           
 36         pkgs,                                                                                                 
 37         system,                                                                                               
 38         ...                                                                                                   
 39       }: {                                                                                                    
 40         _module.args.pkgs = import nixpkgs {                                                                  
 41           inherit system;                                                                                     
 42           config.allowUnfree = true;                                                                          
 43           config.cudaSupport = true;                                                                          
 44         };

Sorry about the line numbers but you just do this.

In reply to @trexd:matrix.org
Oh I reread your message. Cant you just do nixpkgs.config.cudaSupport = true; in your configuration.nix if you want to globally enable cudaSupport?

In reply to @ss:someonex.net
Wdym by "single input" and by the "closure with packages for a component"?

In reply to @aidalgol:matrix.org
Thread about zluda over on Discourse: https://discourse.nixos.org/t/overlaying-packages-using-cuda-to-use-zluda/45374
They're trying to get cudaPackages to use zluda.

In reply to @evax:matrix.org
Hi I'm trying to get jax with cuda to work in WSL using a flake, but the GPU is never recognized. Torch in the same flake recognizes the device. I've tried setting nixpkgs to both nixos-23.11 and nixos-unstable.

{
  description = "Jax+cuda shell";

  nixConfig = {
    extra-substituters = [
      "https://cuda-maintainers.cachix.org"
    ];
    extra-trusted-public-keys = [
      "cuda-maintainers.cachix.org-1:0dq3bujKpuEPMCX6U4WylrUDZ9JyUG0VpVZa7CNfq5E="
    ];
  };

  inputs = {
    nixpkgs.url = "github:NixOS/nixpkgs/nixos-23.11";
    flake-utils.url = "github:numtide/flake-utils";
  };

  outputs = { self, nixpkgs, flake-utils }:
    flake-utils.lib.eachDefaultSystem (system:
      let
        config = {
          allowUnfree = true;
          cudaSupport = true;
        };
        pkgs = (import nixpkgs { inherit system config; }).pkgs;
        python3 = pkgs.python311;
        deps = ps: with ps; [
          jax
          jaxlib
        ];
        devPython = python3.withPackages(ps: with ps; (deps(ps) ++ [
          ipython
        ]));
      in rec {
        inherit pkgs;
        devShell = pkgs.stdenv.mkDerivation {
          name = "jax-shell";
          buildInputs = [
            devPython
          ];
          shellHook = ''
            export CUDA_PATH=${pkgs.cudatoolkit}
            export LD_LIBRARY_PATH=/usr/lib/wsl/lib:${pkgs.linuxPackages.nvidia_x11}/lib:${pkgs.ncurses5}/lib
            export EXTRA_LDFLAGS="-L/lib -L${pkgs.linuxPackages.nvidia_x11}/lib"
            export EXTRA_CCFLAGS="-I/usr/include"
         '';
        };
        defaultPackage = devShell;
      }
    );
}

export LD_LIBRARY_PATH=/usr/lib/wsl/lib:${pkgs.linuxPackages.nvidia_x11}/lib:${pkgs.ncurses5}/lib

You don't want to reference nvidia_x11 in WSL environments. Even on linux we don't reference it directly, cf. various posts about "impure drivers" on github and discourse

/usr/lib/wsl/lib:

I forget now, is libcuda.so placed directly under this path, or in a subdirectory?