WASM: defer eager Numba gufunc/jit compilation to first use (#930)

quantecon/game_theory/vertex_enumeration.py

@@ -296,25 +296,33 @@ def __init__(self, opponent_player, idx=0, qhull_options=None):
         self.trans_recip = trans_recip
 
 
-@guvectorize(['(i4[:], u8[:])'], '(m)->()', nopython=True, cache=True)
-def _ints_arr_to_bits(ints_arr, out):
+def _ints_arr_to_bits_kernel(ints_arr, out):  # pragma: no cover
+    m = ints_arr.shape[0]
+    out[0] = 0
+    for i in range(m):
+        out[0] |= np.uint64(1) << np.uint64(ints_arr[i])
+
+
+_ints_arr_to_bits_gufunc = None
+
+
+def _ints_arr_to_bits(ints_arr):
     """
     Convert an array of integers representing the set bits into the
     corresponding integer.
 
-    Compiled as a ufunc by Numba's `@guvectorize`: if the input is a
-    2-dim array with shape[0]=K, the function returns a 1-dim array of
-    K converted integers.
+    If the input is a 2-dim array with shape[0]=K, returns a 1-dim array
+    of K converted integers.
 
     Parameters
     ----------
-    ints_arr : ndarray(int32, ndim=1)
+    ints_arr : ndarray(int32, ndim=1 or 2)
         Array of distinct integers from 0, ..., 63.
 
     Returns
     -------
-    np.uint64
-        Integer with set bits represented by the input integers.
+    np.uint64 or ndarray(uint64)
+        Integer(s) with set bits represented by the input integers.
 
     Examples
     --------
@@ -326,10 +334,12 @@ def _ints_arr_to_bits(ints_arr, out):
     array([ 7, 11], dtype=uint64)
 
     """
-    m = ints_arr.shape[0]
-    out[0] = 0
-    for i in range(m):
-        out[0] |= np.uint64(1) << np.uint64(ints_arr[i])
+    global _ints_arr_to_bits_gufunc
+    if _ints_arr_to_bits_gufunc is None:
+        _ints_arr_to_bits_gufunc = guvectorize(
+            ['(i4[:], u8[:])'], '(m)->()', nopython=True, cache=True
+        )(_ints_arr_to_bits_kernel)
+    return _ints_arr_to_bits_gufunc(ints_arr)
 
 
 @jit(nopython=True, cache=True)

quantecon/random/utilities.py

@@ -89,14 +89,26 @@ def _probvec(r, out):  # pragma: no cover
         out[i] = r[i] - r[i-1]
     out[n] = 1 - r[n-1]
 
-_probvec_parallel = guvectorize(
-    ['(f8[:], f8[:])'], '(n), (k)', nopython=True, target='parallel',
-    cache=True
-    )(_probvec)
-_probvec_cpu = guvectorize(
-    ['(f8[:], f8[:])'], '(n), (k)', nopython=True, target='cpu',
-    cache=True
-    )(_probvec)
+_probvec_parallel_gufunc = None
+_probvec_cpu_gufunc = None
+
+
+def _probvec_parallel(r, out):
+    global _probvec_parallel_gufunc
+    if _probvec_parallel_gufunc is None:
+        _probvec_parallel_gufunc = guvectorize(
+            ['(f8[:], f8[:])'], '(n),(k)', nopython=True,
+            target='parallel', cache=True)(_probvec)
+    _probvec_parallel_gufunc(r, out)
+
+
+def _probvec_cpu(r, out):
+    global _probvec_cpu_gufunc
+    if _probvec_cpu_gufunc is None:
+        _probvec_cpu_gufunc = guvectorize(
+            ['(f8[:], f8[:])'], '(n),(k)', nopython=True,
+            target='cpu', cache=True)(_probvec)
+    _probvec_cpu_gufunc(r, out)
 
 
 def sample_without_replacement(n, k, num_trials=None, random_state=None):
@@ -153,10 +165,9 @@ def sample_without_replacement(n, k, num_trials=None, random_state=None):
     return result
 
 
-@guvectorize(['(i8, f8[:], i8[:])'], '(),(k)->(k)', nopython=True, cache=True)
-def _sample_without_replacement(n, r, out):
+def _sample_without_replacement_kernel(n, r, out):  # pragma: no cover
     """
-    Main body of `sample_without_replacement`. To be complied as a ufunc
+    Main body of `sample_without_replacement`. To be compiled as a ufunc
     by guvectorize of Numba.
 
     """
@@ -170,6 +181,18 @@ def _sample_without_replacement(n, r, out):
         pool[idx] = pool[n-j-1]
 
 
+_swr_gufunc = None
+
+
+def _sample_without_replacement(n, r):
+    global _swr_gufunc
+    if _swr_gufunc is None:
+        _swr_gufunc = guvectorize(
+            ['(i8, f8[:], i8[:])'], '(),(k)->(k)', nopython=True, cache=True
+        )(_sample_without_replacement_kernel)
+    return _swr_gufunc(n, r)
+
+
 # Pure python implementation that will run if the JIT compiler is disabled
 def draw(cdf, size=None, rng=None):
     """