diff --git a/tests/test_2799_numba_ufunc_resolution.py b/tests/test_2799_numba_ufunc_resolution.py index 3806710627..ef4cb14cf4 100644 --- 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