From a225c7d3073fb5bdef090a1b7986d18a8d557b1a Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Sat, 5 Oct 2024 12:48:01 +0200
Subject: [PATCH] python312Packages.tinygrad: fix libnvrtc patching

---
 .../python-modules/tinygrad/default.nix            | 14 +++++++++++++-
 .../python-modules/tinygrad/fix-dlopen-cuda.patch  |  2 +-
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/pkgs/development/python-modules/tinygrad/default.nix b/pkgs/development/python-modules/tinygrad/default.nix
index 14fd91a1bb12..8566c0427112 100644
--- a/pkgs/development/python-modules/tinygrad/default.nix
+++ b/pkgs/development/python-modules/tinygrad/default.nix
@@ -36,6 +36,8 @@
   torch,
   tqdm,
   transformers,
+
+  tinygrad,
 }:
 
 buildPythonPackage rec {
@@ -92,7 +94,13 @@ buildPythonPackage rec {
       # pyobjc-framework-metal
     ];
 
-  pythonImportsCheck = [ "tinygrad" ];
+  pythonImportsCheck =
+    [
+      "tinygrad"
+    ]
+    ++ lib.optionals cudaSupport [
+      "tinygrad.runtime.ops_nv"
+    ];
 
   nativeCheckInputs = [
     blobfile
@@ -175,6 +183,10 @@ buildPythonPackage rec {
     "extra/"
   ];
 
+  passthru.tests = {
+    withCuda = tinygrad.override { cudaSupport = true; };
+  };
+
   meta = {
     description = "Simple and powerful neural network framework";
     homepage = "https://github.com/tinygrad/tinygrad";
diff --git a/pkgs/development/python-modules/tinygrad/fix-dlopen-cuda.patch b/pkgs/development/python-modules/tinygrad/fix-dlopen-cuda.patch
index db74eab45bc0..285ca6111569 100644
--- a/pkgs/development/python-modules/tinygrad/fix-dlopen-cuda.patch
+++ b/pkgs/development/python-modules/tinygrad/fix-dlopen-cuda.patch
@@ -43,7 +43,7 @@ index 6af74187..c5a6c6c4 100644
 +    pass
 +if libnvrtc is None:
 +    raise RuntimeError(f"`libnvrtc.so` not found")
-+_libraries['libnvrtc.so'] = ctypes.CDLL(libnvrtc)
++_libraries['libnvrtc.so'] = libnvrtc
  def string_cast(char_pointer, encoding='utf-8', errors='strict'):
      value = ctypes.cast(char_pointer, ctypes.c_char_p).value
      if value is not None and encoding is not None: