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?

In reply to @evax:matrix.org
some finding, under NixOS-WSL, with the option to use windows drivers, a wsl-lib package is created in the nix store linking the contents of /usr/lib/wsl/lib

In reply to @evax:matrix.org
it might just be that jax is expecting cuda12 but the actual version in the system is cuda11