WASM: defer Numba gufunc compilation to first call (#930)

quantecon/game_theory/random.py

@@ -202,7 +202,7 @@ def _random_mixed_actions(out, random_state):
             x[0] = 1
         else:
             r = random_state.random(size=n-1)
-            _probvec_cpu(r, x)
+            _probvec_cpu()(r, x)
     return out
 
 

quantecon/game_theory/vertex_enumeration.py

@@ -111,8 +111,12 @@ def vertex_enumeration_gen(g, qhull_options=None):
         g.players[1-i], idx=i, qhull_options=qhull_options
     ) for i in range(N)]
 
-    labelings_bits_tup = \
-        tuple(_ints_arr_to_bits(brps[i].labelings) for i in range(N))
+    labelings_bits_tup = tuple(
+        _ints_arr_to_bits(
+            brps[i].labelings,
+            np.empty(brps[i].labelings.shape[0], dtype=np.uint64)
+        ) for i in range(N)
+    )
     equations_tup = tuple(brps[i].equations for i in range(N))
     trans_recips = tuple(brps[i].trans_recip for i in range(N))
 
@@ -302,33 +306,37 @@ def __init__(self, opponent_player, idx=0, qhull_options=None):
         self.trans_recip = trans_recip
 
 
-@guvectorize(['(i4[:], u8[:])'], '(m)->()', nopython=True, cache=True)
+@guvectorize('(m)->()', nopython=True, cache=True)
 def _ints_arr_to_bits(ints_arr, out):
     """
     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.
+    Compiled as a dynamic (signature-less) gufunc by Numba's
+    `@guvectorize` on first call: if the input is a 2-dim array with
+    shape[0]=K, the function fills the 1-dim output array `out` with
+    K converted integers. The output array must be supplied by the
+    caller, with dtype np.uint64.
 
     Parameters
     ----------
-    ints_arr : ndarray(int32, ndim=1)
-        Array of distinct integers from 0, ..., 63.
+    ints_arr : ndarray(int, ndim=2)
+        Array of distinct integers from 0, ..., 63 in each row.
+
+    out : ndarray(np.uint64, ndim=1)
+        Output array, of length ints_arr.shape[0].
 
     Returns
     -------
-    np.uint64
-        Integer with set bits represented by the input integers.
+    out : ndarray(np.uint64, ndim=1)
+        Array of integers with set bits represented by the rows of
+        the input.
 
     Examples
     --------
-    >>> ints_arr = np.array([0, 1, 2], dtype=np.int32)
-    >>> _ints_arr_to_bits(ints_arr)
-    np.uint64(7)
-    >>> ints_arr2d = np.array([[0, 1, 2], [3, 0, 1]], dtype=np.int32)
-    >>> _ints_arr_to_bits(ints_arr2d)
+    >>> ints_arr = np.array([[0, 1, 2], [3, 0, 1]], dtype=np.int32)
+    >>> out = np.empty(ints_arr.shape[0], dtype=np.uint64)
+    >>> _ints_arr_to_bits(ints_arr, out)
     array([ 7, 11], dtype=uint64)
 
     """

quantecon/random/utilities.py

@@ -3,6 +3,8 @@
 
 """
 
+import functools
+
 import numpy as np
 from numba import guvectorize, types
 from numba import TypingError
@@ -60,9 +62,9 @@ def probvec(m, k, random_state=None, parallel=True):
 
     # Parse Parallel Option #
     if parallel:
-        _probvec_parallel(r, x)
+        _probvec_parallel()(r, x)
     else:
-        _probvec_cpu(r, x)
+        _probvec_cpu()(r, x)
 
     return x
 
@@ -92,14 +94,41 @@ 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)
+# The two gufuncs of `_probvec` are constructed on first call rather than
+# at import, since an explicit signature makes `guvectorize` compile
+# eagerly. The explicit signature is required by the parallel target, and
+# is retained for the cpu target as well: an explicit-signature gufunc and
+# a dynamic (signature-less) gufunc of the same kernel must not be mixed
+# with `cache=True`, as they share a cache entry that is not compatible.
+
+_PROBVEC_SIGNATURE = ['(f8[:], f8[:])']
+_PROBVEC_LAYOUT = '(n), (k)'
+
+
+@functools.cache
+def _probvec_parallel():
+    """
+    Return the parallel-target gufunc of `_probvec`, compiling it on
+    first call.
+
+    """
+    return guvectorize(
+        _PROBVEC_SIGNATURE, _PROBVEC_LAYOUT, nopython=True,
+        target='parallel', cache=True
+        )(_probvec)
+
+
+@functools.cache
+def _probvec_cpu():
+    """
+    Return the cpu-target gufunc of `_probvec`, compiling it on first
+    call.
+
+    """
+    return guvectorize(
+        _PROBVEC_SIGNATURE, _PROBVEC_LAYOUT, nopython=True,
+        target='cpu', cache=True
+        )(_probvec)
 
 
 def sample_without_replacement(n, k, num_trials=None, random_state=None):
@@ -155,16 +184,18 @@ def sample_without_replacement(n, k, num_trials=None, random_state=None):
 
     random_state = check_random_state(random_state)
     r = random_state.random(size=size)
-    result = _sample_without_replacement(n, r)
+    result = np.empty(size, dtype=np.int64)
+    _sample_without_replacement(n, r, result)
 
     return result
 
 
-@guvectorize(['(i8, f8[:], i8[:])'], '(),(k)->(k)', nopython=True, cache=True)
+@guvectorize('(),(k)->(k)', nopython=True, cache=True)
 def _sample_without_replacement(n, r, out):
     """
-    Main body of `sample_without_replacement`. To be compiled as a ufunc
-    by guvectorize of Numba.
+    Main body of `sample_without_replacement`. Compiled as a dynamic
+    (signature-less) gufunc by Numba's `@guvectorize` on first call;
+    the output array `out` must be supplied by the caller.
 
     """
     k = r.shape[0]