RFC: Replace @generated_jit with @overload in util.numba

Author: oyamad

quantecon/util/numba.py

@@ -3,7 +3,8 @@
 
 """
 import numpy as np
-from numba import jit, generated_jit, types
+from numba import jit, types
+from numba.extending import overload
 try:
     from numba.np.linalg import _LAPACK  # for Numba >= 0.49.0
 except ModuleNotFoundError:
@@ -19,14 +20,19 @@
 }
 
 
-@generated_jit(nopython=True, cache=True)
-def _numba_linalg_solve(a, b):
+def _numba_linalg_solve(a, b):  # pragma: no cover
+    pass
+
+
+@overload(_numba_linalg_solve, jit_options={'cache':True})
+def _numba_linalg_solve_ol(a, b):
     """
     Solve the linear equation ax = b directly calling a Numba internal
     function. The data in `a` and `b` are interpreted in Fortran order,
     and dtype of `a` and `b` must be the same, one of {float32, float64,
     complex64, complex128}. `a` and `b` are modified in place, and the
     solution is stored in `b`. *No error check is made for the inputs.*
+    Only work in a Numba-jitted function.
 
     Parameters
     ----------

quantecon/util/tests/test_numba.py

@@ -8,6 +8,11 @@
 from quantecon.util.numba import _numba_linalg_solve, comb_jit
 
 
+@jit(nopython=True)
+def _numba_linalg_solve_jitted(a, b):
+    return _numba_linalg_solve(a, b)
+
+
 @jit(nopython=True)
 def numba_linalg_solve_orig(a, b):
     return np.linalg.solve(a, b)
@@ -26,7 +31,7 @@ def test_b_1dim(self):
             a = np.asfortranarray(self.a, dtype=dtype)
             b = np.asfortranarray(self.b_1dim, dtype=dtype)
             sol_orig = numba_linalg_solve_orig(a, b)
-            r = _numba_linalg_solve(a, b)
+            r = _numba_linalg_solve_jitted(a, b)
             assert_(r == 0)
             assert_array_equal(b, sol_orig)
 
@@ -35,7 +40,7 @@ def test_b_2dim(self):
             a = np.asfortranarray(self.a, dtype=dtype)
             b = np.asfortranarray(self.b_2dim, dtype=dtype)
             sol_orig = numba_linalg_solve_orig(a, b)
-            r = _numba_linalg_solve(a, b)
+            r = _numba_linalg_solve_jitted(a, b)
             assert_(r == 0)
             assert_array_equal(b, sol_orig)
 
@@ -44,7 +49,7 @@ def test_singular_a(self):
             for dtype in self.dtypes:
                 a = np.asfortranarray(self.a_singular, dtype=dtype)
                 b = np.asfortranarray(b, dtype=dtype)
-                r = _numba_linalg_solve(a, b)
+                r = _numba_linalg_solve_jitted(a, b)
                 assert_(r != 0)