fix: register Numba for test, add ufunc signatures

scikit-hep · Dec 8, 2023 · 5686284 · 5686284
1 parent 9400780
commit 5686284
Showing 1 changed file with 19 additions and 2 deletions.
diff --git a/tests/test_2799_numba_ufunc_resolution.py b/tests/test_2799_numba_ufunc_resolution.py
@@ -15,11 +15,21 @@
 ) >= packaging.version.Version("0.59.0")
 
 
+ak.numba.register_and_check()
+
+
 @pytest.mark.skipif(not NUMBA_HAS_NEP_50, reason="Numba does not have NEP-50 support")
 def test_numba_ufunc_nep_50():
     raise NotImplementedError
 
-    @numba.vectorize
+    @numba.vectorize(
+        [
+            numba.int8(numba.int8, numba.int8),
+            numba.int16(numba.int16, numba.int16),
+            numba.int32(numba.int32, numba.int32),
+            numba.int64(numba.int64, numba.int64),
+        ]
+    )
     def add(x, y):
         return x + y
 
@@ -35,7 +45,14 @@ def add(x, y):
 
 @pytest.mark.skipif(NUMBA_HAS_NEP_50, reason="Numba has NEP-50 support")
 def test_numba_ufunc_legacy():
-    @numba.vectorize
+    @numba.vectorize(
+        [
+            numba.int8(numba.int8, numba.int8),
+            numba.int16(numba.int16, numba.int16),
+            numba.int32(numba.int32, numba.int32),
+            numba.int64(numba.int64, numba.int64),
+        ]
+    )
     def add(x, y):
         return x + y