Install numba-cuda 0.4.0 and configure numba for JIT link

This is to fix googlecolab/colabtools#5081 (comment)
uradora136 · Feb 26, 2025 · 5a0c69e · 5a0c69e
1 parent b3c4b23
commit 5a0c69e
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/GPU_puzzlers.ipynb b/GPU_puzzlers.ipynb
@@ -128,7 +128,8 @@
       ],
       "source": [
         "!pip install -qqq git+https://github.com/chalk-diagrams/planar git+https://github.com/danoneata/chalk@srush-patch-1\n",
-        "!wget -q https://github.com/srush/GPU-Puzzles/raw/main/robot.png https://github.com/srush/GPU-Puzzles/raw/main/lib.py"
+        "!wget -q https://github.com/srush/GPU-Puzzles/raw/main/robot.png https://github.com/srush/GPU-Puzzles/raw/main/lib.py",
+        "!pip install -qqq numba-cuda==0.4.0"
       ]
     },
     {

diff --git a/lib.py b/lib.py
@@ -6,10 +6,13 @@
 from dataclasses import dataclass
 from typing import List, Any
 from collections import Counter
-from numba import cuda
+from numba import cuda, config
 import numba
 import random
 
+config.CUDA_ENABLE_PYNVJITLINK = 1
+config.CUDA_LOW_OCCUPANCY_WARNINGS = 0
+
 @dataclass
 class ScalarHistory:
     last_fn: str