From 595d15288581960ce1a73831d7f957dda2421526 Mon Sep 17 00:00:00 2001
From: Tibo De Peuter <tibo@depeuter.dev>
Date: Sat, 6 Dec 2025 20:21:43 +0100
Subject: [PATCH] feat(python): Add CUDA

---
 python/uv/flake.nix | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

diff --git a/python/uv/flake.nix b/python/uv/flake.nix
index 3d53c0a..7341edb 100644
--- a/python/uv/flake.nix
+++ b/python/uv/flake.nix
@@ -31,6 +31,15 @@
     devShells.default = pkgs.mkShell {
       packages = (with pkgs; [
         python311Full
+
+        # CUDA
+        git gitRepo gnupg autoconf curl
+        procps gnumake util-linux m4 gperf unzip
+        cudatoolkit linuxPackages.nvidia_x11
+        libGLU libGL
+        xorg.libXi xorg.libXmu freeglut
+        xorg.libXext xorg.libX11 xorg.libXv xorg.libXrandr zlib
+        ncurses5 stdenv.cc binutils
       ]) ++ (with pkgs-unstable; [
         uv
       ]) ++ (with nix-jetbrains-plugins.lib."${system}"; [
@@ -38,11 +47,15 @@
         (buildIdeWithPlugins pkgs-unstable.jetbrains "pycharm-professional" pluginList)
       ]);
 
+      # CUDA
+      CUDA_PATH = pkgs.cudatoolkit;
+      # ImportError: libstdc++.so.6: cannot open shared object file: No such file or directory
+      LD_LIBRARY_PATH = "${pkgs.linuxPackages.nvidia_x11}/lib:${pkgs.ncurses5}/lib:${pkgs.libGL}/lib/:${pkgs.stdenv.cc.cc.lib}/lib/:${pkgs.glibc}/lib";
+      EXTRA_LDFLAGS = "-L/lib -L${pkgs.linuxPackages.nvidia_x11}/lib";
+      EXTRA_CCFLAGS = "-I/usr/include";
+
       # Stop uv from downloading Python binaries automatically if needed.
       UV_PYTHON_DOWNLOADS = "never";
-
-      # ImportError: libstdc++.so.6: cannot open shared object file: No such file or directory
-      LD_LIBRARY_PATH = "${pkgs.stdenv.cc.cc.lib}/lib";
     };
   });
 }