From c58c790785dac4c8c7e0e4a243a79781ffee5570 Mon Sep 17 00:00:00 2001 From: Matt McKay Date: Sat, 1 Aug 2026 11:38:12 +1000 Subject: [PATCH 1/4] ENH: Extend random.draw to accept a random number generator `draw` was the only function in `quantecon.random` with no `random_state` argument. It called `np.random.random()` directly in both the pure Python body and the `@overload` implementation, so the generator it used depended on the call site and `np.random.seed` at Python level had no effect on the jitted path. Add `random_state` as a trailing keyword argument, matching the name already used by `probvec`, `sample_without_replacement` and `DiscreteRV.draw`. The pure Python body routes through `check_random_state`, so it accepts None, an int seed, a RandomState or a Generator. The overload takes a narrower contract -- None or a live `np.random.Generator` -- because nopython mode can neither construct a generator from a seed nor represent a RandomState; anything else raises a TypingError at compile time with a message saying what to pass instead. A Generator reproduces the host stream bit for bit in nopython mode and its state is mutated in place, so one generator stays in sync across Python and jitted calls. That is now the recommended way to get reproducible draws from a call site that may be jit-compiled. The change is purely additive. `check_random_state(None)` returns `np.random.mtrand._rand` by identity and is non-consuming, so the Python None path is byte-identical to the previous behaviour, and the jitted None path still uses Numba's internal state, leaving jitted callers that seed with `np.random.seed` unaffected. The jitted sized branch keeps its hand-written searchsorted loop rather than mirroring the Python body's vectorised call: the two agree exactly but the loop measures about 2x faster at size >= 1000, so `test_python_jitted_agree` is what holds the paths together instead. Also raise the numba floor to 0.59.0. The overload needs `types.NumPyRandomGeneratorType`, added in numba 0.56.0, and 0.59.0 is the first release with cp312 wheels, which the classifiers and CI matrix already require. See #916 Co-Authored-By: Claude Opus 5 (1M context) --- pyproject.toml | 2 +- quantecon/random/tests/test_utilities.py | 180 ++++++++++++++++++++++- quantecon/random/utilities.py | 153 ++++++++++++++++--- 3 files changed, 315 insertions(+), 20 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 8797dad4..6fbd5be0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,7 +23,7 @@ keywords = [ dynamic = ["description", "version"] requires-python = ">=3.7" dependencies = [ - 'numba>=0.49.0', + 'numba>=0.59.0', 'numpy>=1.17.0', 'requests', 'scipy>=1.5.0', diff --git a/quantecon/random/tests/test_utilities.py b/quantecon/random/tests/test_utilities.py index 430c7a1c..1ccc13aa 100644 --- a/quantecon/random/tests/test_utilities.py +++ b/quantecon/random/tests/test_utilities.py @@ -5,13 +5,16 @@ --------- probvec sample_without_replacement +draw """ import numbers import numpy as np +import pytest from numpy.testing import (assert_array_equal, assert_allclose, assert_raises, assert_) -from numba import njit +from numba import config, njit +from numba.core.errors import TypingError from quantecon.random import probvec, sample_without_replacement, draw @@ -73,6 +76,39 @@ def draw_jitted(cdf, size=None): return draw(cdf, size) +@njit +def draw_jitted_rs(cdf, size, random_state): + return draw(cdf, size, random_state) + + +@njit +def draw_jitted_rs_kw(cdf, random_state): + return draw(cdf, random_state=random_state) + + +@njit +def draw_jitted_explicit_none(cdf, size): + return draw(cdf, size, None) + + +@njit +def draw_jitted_fwd(cdf, size=None, random_state=None): + # Forwards its own defaults, which reach the overload as types.none + return draw(cdf, size, random_state) + + +@njit +def draw_jitted_seeded(cdf, size, seed): + np.random.seed(seed) + return draw(cdf, size) + + +@njit +def draw_jitted_optional_rs(cdf, size, random_state, flag): + rs = random_state if flag else None + return draw(cdf, size, rs) + + class TestDraw: def setup_method(self): self.pmf = np.array([0.4, 0.1, 0.5]) @@ -109,6 +145,148 @@ def test_lln(self): atol = 1e-2 assert_allclose(pmf_computed, self.pmf, atol=atol) + # random_state: a Generator behaves the same on both paths # + + def test_python_jitted_agree(self): + # The contract that keeps the two implementations in step: for + # the same Generator seed they must return the same values. + for seed in [0, 1234, 20260801]: + for size in [1, 10, 1000]: + out_py = draw(self.cdf, size, np.random.default_rng(seed)) + out_jit = draw_jitted_rs(self.cdf, size, + np.random.default_rng(seed)) + assert_array_equal(out_py, out_jit) + + # Scalar draws are compared as arrays: the Python path + # returns np.int64 and the jitted path a Python int. + out_py = draw(self.cdf, None, np.random.default_rng(seed)) + out_jit = draw_jitted_rs(self.cdf, None, + np.random.default_rng(seed)) + assert_array_equal(np.asarray(out_py), np.asarray(out_jit)) + + def test_generator_reproducible(self): + for size in [None, 10]: + for func in [draw, draw_jitted_rs]: + out0 = func(self.cdf, size, np.random.default_rng(1234)) + out1 = func(self.cdf, size, np.random.default_rng(1234)) + assert_array_equal(np.asarray(out0), np.asarray(out1)) + + def test_generator_state_advances(self): + # The Generator must be mutated in place by the jitted call, not + # copied at the boundary, so that successive calls continue the + # stream rather than restarting it. + rng = np.random.default_rng(1234) + parts = [draw_jitted_rs(self.cdf, 5, rng), + draw_jitted_rs(self.cdf, 5, rng), + draw(self.cdf, 5, rng)] + expected = draw(self.cdf, 15, np.random.default_rng(1234)) + assert_array_equal(np.concatenate(parts), expected) + + def test_consumes_exactly_one_draw_per_variate(self): + # An extra variate taken at the *end* of a call is invisible to + # the value comparisons above, but it desynchronises a shared + # Generator for every later consumer. Complements + # test_generator_state_advances, which covers the in-place + # mutation half of the same contract. + for func in [draw, draw_jitted_rs]: + for size in [None, 1, 10]: + n = 1 if size is None else size + rng = np.random.default_rng(1234) + func(self.cdf, size, rng) + ref = np.random.default_rng(1234) + ref.random(n) + assert_array_equal(rng.random(4), ref.random(4)) + + def test_generator_keyword_form_in_jit(self): + out_py = draw(self.cdf, random_state=np.random.default_rng(1234)) + out_jit = draw_jitted_rs_kw(self.cdf, np.random.default_rng(1234)) + assert_array_equal(np.asarray(out_py), np.asarray(out_jit)) + + # random_state: the Python path keeps check_random_state's breadth # + + def test_int_seed_python_path(self): + for size in [None, 10]: + out_seed = draw(self.cdf, size, 1234) + out_rs = draw(self.cdf, size, np.random.RandomState(1234)) + assert_array_equal(np.asarray(out_seed), np.asarray(out_rs)) + + def test_int_seed_and_generator_differ(self): + # Documented in Notes: same integer, different stream. + out_seed = draw(self.cdf, 10, 1234) + out_gen = draw(self.cdf, 10, np.random.default_rng(1234)) + assert_(not np.array_equal(out_seed, out_gen)) + + # random_state=None: unchanged behaviour on both paths # + + def test_none_matches_legacy_global(self): + for size in [None, 10]: + np.random.seed(99) + out = draw(self.cdf, size) + np.random.seed(99) + r = np.random.random(size) if size is not None \ + else np.random.random() + expected = np.searchsorted(self.cdf, r, side='right') + assert_array_equal(np.asarray(out), np.asarray(expected)) + + def test_none_spellings_all_compile(self): + # Omitted, explicitly None, and forwarded from a jitted caller's + # own default are distinct numba types, and each must reach the + # np.random branch of the overload. + size = 10 + for out in [draw_jitted(self.cdf, size), + draw_jitted_explicit_none(self.cdf, size), + draw_jitted_fwd(self.cdf, size), + draw_jitted_fwd(self.cdf, size, None)]: + assert_(out.shape == (size,)) + assert_(np.isin(out, range(self.n)).all()) + + for out in [draw_jitted(self.cdf), draw_jitted_fwd(self.cdf)]: + assert_(out in range(self.n)) + + def test_jitted_np_random_seed_still_reproducible(self): + out0 = draw_jitted_seeded(self.cdf, 10, 1234) + out1 = draw_jitted_seeded(self.cdf, 10, 1234) + assert_array_equal(out0, out1) + + # random_state: rejections in nopython mode # + + @pytest.mark.skipif(config.DISABLE_JIT, + reason='requires nopython compilation') + def test_int_seed_raises_in_jit(self): + assert_raises(TypingError, draw_jitted_rs, self.cdf, 10, 1234) + try: + draw_jitted_rs(self.cdf, 10, 1234) + except TypingError as e: + # Numba nests the message in its report of the candidate + # implementations it rejected. Assert on tokens that can + # only come from draw's own message, not from the caller's + # source line that numba echoes alongside it. + msg = str(e) + assert_('quantecon.random.draw' in msg) + assert_('np.random.default_rng' in msg) + + @pytest.mark.skipif(config.DISABLE_JIT, + reason='requires nopython compilation') + def test_randomstate_raises_in_jit(self): + # A RandomState cannot be typed at all, so it is rejected during + # argument typing and the overload never runs. Assert only the + # exception type, never the message. + assert_raises(TypingError, draw_jitted_rs, self.cdf, 10, + np.random.RandomState(1234)) + + @pytest.mark.skipif(config.DISABLE_JIT, + reason='requires nopython compilation') + def test_optional_generator_raises_in_jit(self): + assert_raises(TypingError, draw_jitted_optional_rs, self.cdf, 10, + np.random.default_rng(1234), True) + try: + draw_jitted_optional_rs(self.cdf, 10, + np.random.default_rng(1234), True) + except TypingError as e: + # A token unique to the Optional hint, so that this branch + # cannot be satisfied by the int-seed message. + assert_('could not prove' in str(e)) + @njit def draw_jitted_w_o_size(n): diff --git a/quantecon/random/utilities.py b/quantecon/random/utilities.py index bd309b06..4be48577 100644 --- a/quantecon/random/utilities.py +++ b/quantecon/random/utilities.py @@ -5,6 +5,7 @@ import numpy as np from numba import guvectorize, types +from numba.core.errors import TypingError from numba.extending import overload from ..util import check_random_state @@ -170,7 +171,7 @@ def _sample_without_replacement(n, r, out): # Pure python implementation that will run if the JIT compiler is disabled -def draw(cdf, size=None): +def draw(cdf, size=None, random_state=None): """ Generate a random sample according to the cumulative distribution given by `cdf`. Jit-complied by Numba in nopython mode. @@ -185,40 +186,156 @@ def draw(cdf, size=None): `size` independent draws is returned; otherwise, a single draw is returned as a scalar. + random_state : int or np.random.RandomState/Generator, optional + Random seed (integer) or np.random.RandomState or Generator + instance to set the initial state of the random number generator + for reproducibility. If None, the global random state described + in Notes is used. Called from within a jit-compiled function, + only a Generator or None is accepted. + Returns ------- scalar(int) or ndarray(int, ndim=1) + Notes + ----- + An `np.random.Generator` is the only value of `random_state` that + behaves identically on both paths: called from Python and called + from within a jit-compiled function it consumes the same generator + and returns the same draws. Pass one whenever the call site may be + jit-compiled. + + Numba's nopython mode can neither construct a generator nor + represent `np.random.RandomState`, so an integer seed and a + `RandomState` are accepted only on the pure Python path; supplying + either from within a jit-compiled function raises a compile-time + `numba.TypingError`. Note also that a `Generator` and a + `RandomState` seeded with the same integer produce different + streams, so ``draw(cdf, n, 1234)`` and + ``draw(cdf, n, np.random.default_rng(1234))`` are each reproducible + but are not equal to each other. + + If `random_state` is None the draws come from a global random state + that depends on the call site: NumPy's legacy global random state + from Python, seeded by `np.random.seed`, and Numba's own internal + random state from within a jit-compiled function, seeded by calling + `np.random.seed` inside the jit-compiled function. Both are Mersenne + Twister and give the same stream for the same seed, but they advance + independently, so seeding one has no effect on the other. + + A single `Generator` must not be shared across the iterations of a + `numba.prange` loop. Its state is updated in place, so concurrent + iterations race and can consume the same underlying random number; + the draws are then neither reproducible nor independent. Give each + iteration its own generator, spawned with `Generator.spawn` outside + the jit-compiled function -- `spawn` itself cannot be called in + nopython mode. + Examples -------- + >>> import numpy as np + >>> import quantecon as qe >>> cdf = np.cumsum([0.4, 0.6]) - >>> qe.random.draw(cdf) - 1 - >>> qe.random.draw(cdf, 10) - array([1, 0, 1, 0, 1, 0, 0, 0, 1, 0]) + >>> qe.random.draw(cdf, 10, random_state=1234) + array([0, 1, 1, 1, 1, 0, 0, 1, 1, 1]) + >>> int(qe.random.draw(cdf, random_state=1234)) + 0 + + A `Generator` gives the same draws from Python and from within a + jit-compiled function: + + >>> from numba import njit + >>> @njit + ... def draw_jitted(cdf, size, random_state): + ... return qe.random.draw(cdf, size, random_state) + >>> qe.random.draw(cdf, 5, np.random.default_rng(1234)) + array([1, 0, 1, 0, 0]) + >>> draw_jitted(cdf, 5, np.random.default_rng(1234)) + array([1, 0, 1, 0, 0]) """ + random_state = check_random_state(random_state) if isinstance(size, int): - rs = np.random.random(size) + rs = random_state.random(size) out = np.searchsorted(cdf, rs, side='right') return out else: - r = np.random.random() + r = random_state.random() return np.searchsorted(cdf, r, side='right') +def _is_no_random_state(numba_type): + """ + Return True if `numba_type`, as seen from inside the `draw` + overload, means that `random_state` was not supplied. + + Numba spells "no value" in more than one way depending on the call + site, and each must be recognised: an omitted argument, as in + ``draw(cdf, 10)``, arrives as the *Python* object `None`, while an + explicit `None`, as in ``draw(cdf, 10, None)``, and a `None` default + forwarded by a jit-compiled caller both arrive as + `numba.types.none`. `numba.types.Omitted` is not currently observed + in an `@overload` body but is handled for the benefit of other Numba + typing templates. + + """ + return (numba_type is None or + isinstance(numba_type, types.NoneType) or + (isinstance(numba_type, types.Omitted) and + numba_type.value is None)) + + # Overload for the `draw` function +# +# The implementations below must return the same values as the pure +# Python body above: the same random numbers drawn in the same order, +# and the same `searchsorted` result. The sized branches keep the +# hand-written loop rather than the array `np.searchsorted` used in the +# Python body -- Numba supports both and they agree exactly, but the +# loop is measurably faster. `TestDraw.test_python_jitted_agree` is what +# holds the two paths together. @overload(draw) -def ol_draw(cdf, size=None): - if isinstance(size, types.Integer): - def draw_impl(cdf, size=None): - rs = np.random.random(size) - out = np.empty(size, dtype=np.int_) - for i in range(size): - out[i] = np.searchsorted(cdf, rs[i], side='right') - return out +def ol_draw(cdf, size=None, random_state=None): + if isinstance(random_state, types.NumPyRandomGeneratorType): + if isinstance(size, types.Integer): + def draw_impl(cdf, size=None, random_state=None): + rs = random_state.random(size) + out = np.empty(size, dtype=np.int_) + for i in range(size): + out[i] = np.searchsorted(cdf, rs[i], side='right') + return out + else: + def draw_impl(cdf, size=None, random_state=None): + r = random_state.random() + return np.searchsorted(cdf, r, side='right') + elif _is_no_random_state(random_state): + if isinstance(size, types.Integer): + def draw_impl(cdf, size=None, random_state=None): + rs = np.random.random(size) + out = np.empty(size, dtype=np.int_) + for i in range(size): + out[i] = np.searchsorted(cdf, rs[i], side='right') + return out + else: + def draw_impl(cdf, size=None, random_state=None): + r = np.random.random() + return np.searchsorted(cdf, r, side='right') else: - def draw_impl(cdf, size=None): - r = np.random.random() - return np.searchsorted(cdf, r, side='right') + if (isinstance(random_state, types.Optional) and + isinstance(random_state.type, + types.NumPyRandomGeneratorType)): + hint = ('Numba could not prove that `random_state` is a ' + 'Generator rather than None at this call site; hoist ' + 'the None case out of the branch that reaches `draw`.') + else: + hint = ('Integer seeds and np.random.RandomState are accepted ' + 'only on the pure Python path, because nopython mode ' + 'cannot construct a generator. Build one outside the ' + 'jit-compiled function with np.random.default_rng(seed) ' + 'and pass that in.') + raise TypingError( + 'quantecon.random.draw: `random_state` must be None or an ' + 'np.random.Generator when `draw` is called from a jit-compiled ' + f'function; got {random_state}. {hint}' + ) return draw_impl From 59cf867671c6deb540ee18e8f71d8d8f8fb7cf37 Mon Sep 17 00:00:00 2001 From: Matt McKay Date: Fri, 14 Aug 2026 15:24:12 +1000 Subject: [PATCH 2/4] ENH: Narrow draw's RNG argument to rng=None|Generator Adopts the design settled in the PR discussion (#917): the new argument is named rng, and both the Python and jitted paths accept only None or np.random.Generator. Integer seeds and np.random.RandomState now raise TypeError from Python, mirroring the compile-time TypingError from a jit-compiled caller, with the same guidance to build a generator with np.random.default_rng(seed). This is deliberately stricter than SPEC 7: the default_rng(rng) normalization it recommends is impossible in nopython mode, and accepting seeds on only one path would let the two paths disagree. Also, per review comments: fix the "Jit-complied" typo in the draw docstring, correct the stale test-module docstring (util/random.py -> random/utilities.py), and bump requires-python to >=3.12 to match the declared classifiers and the numba>=0.59.0 floor. Co-Authored-By: Claude Fable 5 --- pyproject.toml | 2 +- quantecon/random/tests/test_utilities.py | 102 ++++++++---------- quantecon/random/utilities.py | 128 ++++++++++++----------- 3 files changed, 112 insertions(+), 120 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 6fbd5be0..d18cfadc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,7 +21,7 @@ keywords = [ 'economics' ] dynamic = ["description", "version"] -requires-python = ">=3.7" +requires-python = ">=3.12" dependencies = [ 'numba>=0.59.0', 'numpy>=1.17.0', diff --git a/quantecon/random/tests/test_utilities.py b/quantecon/random/tests/test_utilities.py index 1ccc13aa..79248023 100644 --- a/quantecon/random/tests/test_utilities.py +++ b/quantecon/random/tests/test_utilities.py @@ -1,5 +1,5 @@ """ -Tests for util/random.py +Tests for random/utilities.py Functions --------- @@ -77,13 +77,13 @@ def draw_jitted(cdf, size=None): @njit -def draw_jitted_rs(cdf, size, random_state): - return draw(cdf, size, random_state) +def draw_jitted_rng(cdf, size, rng): + return draw(cdf, size, rng) @njit -def draw_jitted_rs_kw(cdf, random_state): - return draw(cdf, random_state=random_state) +def draw_jitted_rng_kw(cdf, rng): + return draw(cdf, rng=rng) @njit @@ -92,9 +92,9 @@ def draw_jitted_explicit_none(cdf, size): @njit -def draw_jitted_fwd(cdf, size=None, random_state=None): +def draw_jitted_fwd(cdf, size=None, rng=None): # Forwards its own defaults, which reach the overload as types.none - return draw(cdf, size, random_state) + return draw(cdf, size, rng) @njit @@ -104,9 +104,9 @@ def draw_jitted_seeded(cdf, size, seed): @njit -def draw_jitted_optional_rs(cdf, size, random_state, flag): - rs = random_state if flag else None - return draw(cdf, size, rs) +def draw_jitted_optional_rng(cdf, size, rng, flag): + rng_or_none = rng if flag else None + return draw(cdf, size, rng_or_none) class TestDraw: @@ -145,7 +145,7 @@ def test_lln(self): atol = 1e-2 assert_allclose(pmf_computed, self.pmf, atol=atol) - # random_state: a Generator behaves the same on both paths # + # rng: a Generator behaves the same on both paths # def test_python_jitted_agree(self): # The contract that keeps the two implementations in step: for @@ -153,20 +153,20 @@ def test_python_jitted_agree(self): for seed in [0, 1234, 20260801]: for size in [1, 10, 1000]: out_py = draw(self.cdf, size, np.random.default_rng(seed)) - out_jit = draw_jitted_rs(self.cdf, size, - np.random.default_rng(seed)) + out_jit = draw_jitted_rng(self.cdf, size, + np.random.default_rng(seed)) assert_array_equal(out_py, out_jit) # Scalar draws are compared as arrays: the Python path # returns np.int64 and the jitted path a Python int. out_py = draw(self.cdf, None, np.random.default_rng(seed)) - out_jit = draw_jitted_rs(self.cdf, None, - np.random.default_rng(seed)) + out_jit = draw_jitted_rng(self.cdf, None, + np.random.default_rng(seed)) assert_array_equal(np.asarray(out_py), np.asarray(out_jit)) def test_generator_reproducible(self): for size in [None, 10]: - for func in [draw, draw_jitted_rs]: + for func in [draw, draw_jitted_rng]: out0 = func(self.cdf, size, np.random.default_rng(1234)) out1 = func(self.cdf, size, np.random.default_rng(1234)) assert_array_equal(np.asarray(out0), np.asarray(out1)) @@ -176,8 +176,8 @@ def test_generator_state_advances(self): # copied at the boundary, so that successive calls continue the # stream rather than restarting it. rng = np.random.default_rng(1234) - parts = [draw_jitted_rs(self.cdf, 5, rng), - draw_jitted_rs(self.cdf, 5, rng), + parts = [draw_jitted_rng(self.cdf, 5, rng), + draw_jitted_rng(self.cdf, 5, rng), draw(self.cdf, 5, rng)] expected = draw(self.cdf, 15, np.random.default_rng(1234)) assert_array_equal(np.concatenate(parts), expected) @@ -188,7 +188,7 @@ def test_consumes_exactly_one_draw_per_variate(self): # Generator for every later consumer. Complements # test_generator_state_advances, which covers the in-place # mutation half of the same contract. - for func in [draw, draw_jitted_rs]: + for func in [draw, draw_jitted_rng]: for size in [None, 1, 10]: n = 1 if size is None else size rng = np.random.default_rng(1234) @@ -198,25 +198,11 @@ def test_consumes_exactly_one_draw_per_variate(self): assert_array_equal(rng.random(4), ref.random(4)) def test_generator_keyword_form_in_jit(self): - out_py = draw(self.cdf, random_state=np.random.default_rng(1234)) - out_jit = draw_jitted_rs_kw(self.cdf, np.random.default_rng(1234)) + out_py = draw(self.cdf, rng=np.random.default_rng(1234)) + out_jit = draw_jitted_rng_kw(self.cdf, np.random.default_rng(1234)) assert_array_equal(np.asarray(out_py), np.asarray(out_jit)) - # random_state: the Python path keeps check_random_state's breadth # - - def test_int_seed_python_path(self): - for size in [None, 10]: - out_seed = draw(self.cdf, size, 1234) - out_rs = draw(self.cdf, size, np.random.RandomState(1234)) - assert_array_equal(np.asarray(out_seed), np.asarray(out_rs)) - - def test_int_seed_and_generator_differ(self): - # Documented in Notes: same integer, different stream. - out_seed = draw(self.cdf, 10, 1234) - out_gen = draw(self.cdf, 10, np.random.default_rng(1234)) - assert_(not np.array_equal(out_seed, out_gen)) - - # random_state=None: unchanged behaviour on both paths # + # rng=None: unchanged behaviour on both paths # def test_none_matches_legacy_global(self): for size in [None, 10]: @@ -248,40 +234,42 @@ def test_jitted_np_random_seed_still_reproducible(self): out1 = draw_jitted_seeded(self.cdf, 10, 1234) assert_array_equal(out0, out1) - # random_state: rejections in nopython mode # - - @pytest.mark.skipif(config.DISABLE_JIT, - reason='requires nopython compilation') - def test_int_seed_raises_in_jit(self): - assert_raises(TypingError, draw_jitted_rs, self.cdf, 10, 1234) - try: - draw_jitted_rs(self.cdf, 10, 1234) - except TypingError as e: + # rng: rejections on both paths # + + def test_int_seed_raises(self): + # Stricter than SPEC 7 by design: nopython mode cannot construct + # a generator from a seed, so the Python path rejects seeds too + # rather than accept a value the jitted path never can. From + # Python this is a TypeError; from a jitted caller a + # compile-time TypingError (or the same TypeError when + # NUMBA_DISABLE_JIT routes the wrapper to the Python body). + for func in [draw, draw_jitted_rng]: + with pytest.raises((TypeError, TypingError)) as excinfo: + func(self.cdf, 10, 1234) # Numba nests the message in its report of the candidate # implementations it rejected. Assert on tokens that can # only come from draw's own message, not from the caller's # source line that numba echoes alongside it. - msg = str(e) + msg = str(excinfo.value) assert_('quantecon.random.draw' in msg) assert_('np.random.default_rng' in msg) - @pytest.mark.skipif(config.DISABLE_JIT, - reason='requires nopython compilation') - def test_randomstate_raises_in_jit(self): - # A RandomState cannot be typed at all, so it is rejected during - # argument typing and the overload never runs. Assert only the - # exception type, never the message. - assert_raises(TypingError, draw_jitted_rs, self.cdf, 10, - np.random.RandomState(1234)) + def test_randomstate_raises(self): + # In nopython mode a RandomState cannot be typed at all, so it + # is rejected during argument typing and the overload never + # runs. Assert only the exception type, never the message. + for func in [draw, draw_jitted_rng]: + assert_raises((TypeError, TypingError), func, self.cdf, 10, + np.random.RandomState(1234)) @pytest.mark.skipif(config.DISABLE_JIT, reason='requires nopython compilation') def test_optional_generator_raises_in_jit(self): - assert_raises(TypingError, draw_jitted_optional_rs, self.cdf, 10, + assert_raises(TypingError, draw_jitted_optional_rng, self.cdf, 10, np.random.default_rng(1234), True) try: - draw_jitted_optional_rs(self.cdf, 10, - np.random.default_rng(1234), True) + draw_jitted_optional_rng(self.cdf, 10, + np.random.default_rng(1234), True) except TypingError as e: # A token unique to the Optional hint, so that this branch # cannot be satisfied by the int-seed message. diff --git a/quantecon/random/utilities.py b/quantecon/random/utilities.py index 4be48577..c6502764 100644 --- a/quantecon/random/utilities.py +++ b/quantecon/random/utilities.py @@ -171,10 +171,10 @@ def _sample_without_replacement(n, r, out): # Pure python implementation that will run if the JIT compiler is disabled -def draw(cdf, size=None, random_state=None): +def draw(cdf, size=None, rng=None): """ Generate a random sample according to the cumulative distribution - given by `cdf`. Jit-complied by Numba in nopython mode. + given by `cdf`. JIT-compiled by Numba in nopython mode. Parameters ---------- @@ -186,12 +186,11 @@ def draw(cdf, size=None, random_state=None): `size` independent draws is returned; otherwise, a single draw is returned as a scalar. - random_state : int or np.random.RandomState/Generator, optional - Random seed (integer) or np.random.RandomState or Generator - instance to set the initial state of the random number generator - for reproducibility. If None, the global random state described - in Notes is used. Called from within a jit-compiled function, - only a Generator or None is accepted. + rng : np.random.Generator, optional(default=None) + Random number generator to draw from. Must be a + `np.random.Generator` or None; in particular, integer seeds and + `np.random.RandomState` are not accepted. If None, the global + random state described in Notes is used. Returns ------- @@ -199,26 +198,25 @@ def draw(cdf, size=None, random_state=None): Notes ----- - An `np.random.Generator` is the only value of `random_state` that - behaves identically on both paths: called from Python and called - from within a jit-compiled function it consumes the same generator - and returns the same draws. Pass one whenever the call site may be - jit-compiled. - - Numba's nopython mode can neither construct a generator nor - represent `np.random.RandomState`, so an integer seed and a - `RandomState` are accepted only on the pure Python path; supplying - either from within a jit-compiled function raises a compile-time - `numba.TypingError`. Note also that a `Generator` and a - `RandomState` seeded with the same integer produce different - streams, so ``draw(cdf, n, 1234)`` and - ``draw(cdf, n, np.random.default_rng(1234))`` are each reproducible - but are not equal to each other. - - If `random_state` is None the draws come from a global random state - that depends on the call site: NumPy's legacy global random state - from Python, seeded by `np.random.seed`, and Numba's own internal - random state from within a jit-compiled function, seeded by calling + `draw` is intended primarily for use inside jit-compiled functions. + A `np.random.Generator` passed as `rng` behaves identically on both + paths: called from Python and called from within a jit-compiled + function it consumes the same generator and returns the same draws, + and the generator's state advances in place, so a single generator + can be shared between Python and jitted call sites and stays in + sync. Pass one whenever the draws should be reproducible. + + `rng` accepts only None or a `Generator` because Numba's nopython + mode can neither construct a generator from a seed nor represent + `np.random.RandomState`. Anything else raises `TypeError` from + Python and a compile-time `numba.TypingError` from a jit-compiled + caller. Build a generator outside any jit-compiled function with + `np.random.default_rng(seed)` and pass it in. + + If `rng` is None the draws come from a global random state that + depends on the call site: NumPy's legacy global random state from + Python, seeded by `np.random.seed`, and Numba's own internal random + state from within a jit-compiled function, seeded by calling `np.random.seed` inside the jit-compiled function. Both are Mersenne Twister and give the same stream for the same seed, but they advance independently, so seeding one has no effect on the other. @@ -236,38 +234,46 @@ def draw(cdf, size=None, random_state=None): >>> import numpy as np >>> import quantecon as qe >>> cdf = np.cumsum([0.4, 0.6]) - >>> qe.random.draw(cdf, 10, random_state=1234) - array([0, 1, 1, 1, 1, 0, 0, 1, 1, 1]) - >>> int(qe.random.draw(cdf, random_state=1234)) - 0 + >>> rng = np.random.default_rng(1234) + >>> qe.random.draw(cdf, 10, rng=rng) + array([1, 0, 1, 0, 0, 0, 0, 0, 1, 0]) + >>> int(qe.random.draw(cdf, rng=np.random.default_rng(1234))) + 1 A `Generator` gives the same draws from Python and from within a jit-compiled function: >>> from numba import njit >>> @njit - ... def draw_jitted(cdf, size, random_state): - ... return qe.random.draw(cdf, size, random_state) + ... def draw_jitted(cdf, size, rng): + ... return qe.random.draw(cdf, size, rng) >>> qe.random.draw(cdf, 5, np.random.default_rng(1234)) array([1, 0, 1, 0, 0]) >>> draw_jitted(cdf, 5, np.random.default_rng(1234)) array([1, 0, 1, 0, 0]) """ - random_state = check_random_state(random_state) + if rng is None: + rng = np.random + elif not isinstance(rng, np.random.Generator): + raise TypeError( + 'quantecon.random.draw: `rng` must be None or an ' + f'np.random.Generator; got {rng!r}. Build a generator with ' + 'np.random.default_rng(seed) and pass that in.' + ) if isinstance(size, int): - rs = random_state.random(size) + rs = rng.random(size) out = np.searchsorted(cdf, rs, side='right') return out else: - r = random_state.random() + r = rng.random() return np.searchsorted(cdf, r, side='right') -def _is_no_random_state(numba_type): +def _is_no_rng(numba_type): """ Return True if `numba_type`, as seen from inside the `draw` - overload, means that `random_state` was not supplied. + overload, means that `rng` was not supplied. Numba spells "no value" in more than one way depending on the call site, and each must be recognised: an omitted argument, as in @@ -295,47 +301,45 @@ def _is_no_random_state(numba_type): # loop is measurably faster. `TestDraw.test_python_jitted_agree` is what # holds the two paths together. @overload(draw) -def ol_draw(cdf, size=None, random_state=None): - if isinstance(random_state, types.NumPyRandomGeneratorType): +def ol_draw(cdf, size=None, rng=None): + if isinstance(rng, types.NumPyRandomGeneratorType): if isinstance(size, types.Integer): - def draw_impl(cdf, size=None, random_state=None): - rs = random_state.random(size) + def draw_impl(cdf, size=None, rng=None): + rs = rng.random(size) out = np.empty(size, dtype=np.int_) for i in range(size): out[i] = np.searchsorted(cdf, rs[i], side='right') return out else: - def draw_impl(cdf, size=None, random_state=None): - r = random_state.random() + def draw_impl(cdf, size=None, rng=None): + r = rng.random() return np.searchsorted(cdf, r, side='right') - elif _is_no_random_state(random_state): + elif _is_no_rng(rng): if isinstance(size, types.Integer): - def draw_impl(cdf, size=None, random_state=None): + def draw_impl(cdf, size=None, rng=None): rs = np.random.random(size) out = np.empty(size, dtype=np.int_) for i in range(size): out[i] = np.searchsorted(cdf, rs[i], side='right') return out else: - def draw_impl(cdf, size=None, random_state=None): + def draw_impl(cdf, size=None, rng=None): r = np.random.random() return np.searchsorted(cdf, r, side='right') else: - if (isinstance(random_state, types.Optional) and - isinstance(random_state.type, - types.NumPyRandomGeneratorType)): - hint = ('Numba could not prove that `random_state` is a ' - 'Generator rather than None at this call site; hoist ' - 'the None case out of the branch that reaches `draw`.') + if (isinstance(rng, types.Optional) and + isinstance(rng.type, types.NumPyRandomGeneratorType)): + hint = ('Numba could not prove that `rng` is a Generator ' + 'rather than None at this call site; hoist the None ' + 'case out of the branch that reaches `draw`.') else: - hint = ('Integer seeds and np.random.RandomState are accepted ' - 'only on the pure Python path, because nopython mode ' - 'cannot construct a generator. Build one outside the ' - 'jit-compiled function with np.random.default_rng(seed) ' - 'and pass that in.') + hint = ('Integer seeds and np.random.RandomState are not ' + 'accepted, because nopython mode cannot construct a ' + 'generator. Build one outside the jit-compiled ' + 'function with np.random.default_rng(seed) and pass ' + 'that in.') raise TypingError( - 'quantecon.random.draw: `random_state` must be None or an ' - 'np.random.Generator when `draw` is called from a jit-compiled ' - f'function; got {random_state}. {hint}' + 'quantecon.random.draw: `rng` must be None or an ' + f'np.random.Generator; got {rng}. {hint}' ) return draw_impl From edcd4c2d5ae57e505e43086770d693004d9eb00b Mon Sep 17 00:00:00 2001 From: Matt McKay Date: Sat, 15 Aug 2026 17:28:18 +1000 Subject: [PATCH 3/4] MAINT: leave draw's Python body unchecked per review The pure-Python body of `draw` no longer validates `rng`: the None | Generator contract is enforced only by the overload's compile-time TypingError, and Python-path behaviour for other inputs is unspecified. The two Python-path rejection tests are removed and the remaining three are skipped under NUMBA_DISABLE_JIT again. Also: revert the requires-python bump (deferred to #869), lower the numba floor to 0.56.0 (the minimum this feature needs), import TypingError from numba's top level, and show the bare np.int64 repr in the doctest example. Co-Authored-By: Claude Fable 5 --- pyproject.toml | 4 +- quantecon/random/tests/test_utilities.py | 51 +++++++++++++----------- quantecon/random/utilities.py | 22 +++++----- 3 files changed, 38 insertions(+), 39 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index d18cfadc..d24fd0d5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,9 +21,9 @@ keywords = [ 'economics' ] dynamic = ["description", "version"] -requires-python = ">=3.12" +requires-python = ">=3.7" dependencies = [ - 'numba>=0.59.0', + 'numba>=0.56.0', 'numpy>=1.17.0', 'requests', 'scipy>=1.5.0', diff --git a/quantecon/random/tests/test_utilities.py b/quantecon/random/tests/test_utilities.py index 79248023..30020a90 100644 --- a/quantecon/random/tests/test_utilities.py +++ b/quantecon/random/tests/test_utilities.py @@ -13,8 +13,7 @@ import pytest from numpy.testing import (assert_array_equal, assert_allclose, assert_raises, assert_) -from numba import config, njit -from numba.core.errors import TypingError +from numba import config, njit, TypingError from quantecon.random import probvec, sample_without_replacement, draw @@ -234,33 +233,37 @@ def test_jitted_np_random_seed_still_reproducible(self): out1 = draw_jitted_seeded(self.cdf, 10, 1234) assert_array_equal(out0, out1) - # rng: rejections on both paths # + # rng: compile-time rejections, jitted path only # + # + # The None | Generator contract is enforced only by the overload: + # the pure-Python body does no input checking and its behaviour for + # other inputs is unspecified, so nothing here calls `draw` + # directly, and all three tests are skipped when NUMBA_DISABLE_JIT + # routes the jitted wrappers to the Python body. - def test_int_seed_raises(self): + @pytest.mark.skipif(config.DISABLE_JIT, + reason='requires nopython compilation') + def test_int_seed_raises_in_jit(self): # Stricter than SPEC 7 by design: nopython mode cannot construct - # a generator from a seed, so the Python path rejects seeds too - # rather than accept a value the jitted path never can. From - # Python this is a TypeError; from a jitted caller a - # compile-time TypingError (or the same TypeError when - # NUMBA_DISABLE_JIT routes the wrapper to the Python body). - for func in [draw, draw_jitted_rng]: - with pytest.raises((TypeError, TypingError)) as excinfo: - func(self.cdf, 10, 1234) - # Numba nests the message in its report of the candidate - # implementations it rejected. Assert on tokens that can - # only come from draw's own message, not from the caller's - # source line that numba echoes alongside it. - msg = str(excinfo.value) - assert_('quantecon.random.draw' in msg) - assert_('np.random.default_rng' in msg) - - def test_randomstate_raises(self): + # a generator from a seed. + with pytest.raises(TypingError) as excinfo: + draw_jitted_rng(self.cdf, 10, 1234) + # Numba nests the message in its report of the candidate + # implementations it rejected. Assert on tokens that can + # only come from draw's own message, not from the caller's + # source line that numba echoes alongside it. + msg = str(excinfo.value) + assert_('quantecon.random.draw' in msg) + assert_('np.random.default_rng' in msg) + + @pytest.mark.skipif(config.DISABLE_JIT, + reason='requires nopython compilation') + def test_randomstate_raises_in_jit(self): # In nopython mode a RandomState cannot be typed at all, so it # is rejected during argument typing and the overload never # runs. Assert only the exception type, never the message. - for func in [draw, draw_jitted_rng]: - assert_raises((TypeError, TypingError), func, self.cdf, 10, - np.random.RandomState(1234)) + assert_raises(TypingError, draw_jitted_rng, self.cdf, 10, + np.random.RandomState(1234)) @pytest.mark.skipif(config.DISABLE_JIT, reason='requires nopython compilation') diff --git a/quantecon/random/utilities.py b/quantecon/random/utilities.py index c6502764..1906369b 100644 --- a/quantecon/random/utilities.py +++ b/quantecon/random/utilities.py @@ -5,7 +5,7 @@ import numpy as np from numba import guvectorize, types -from numba.core.errors import TypingError +from numba import TypingError from numba.extending import overload from ..util import check_random_state @@ -208,10 +208,12 @@ def draw(cdf, size=None, rng=None): `rng` accepts only None or a `Generator` because Numba's nopython mode can neither construct a generator from a seed nor represent - `np.random.RandomState`. Anything else raises `TypeError` from - Python and a compile-time `numba.TypingError` from a jit-compiled - caller. Build a generator outside any jit-compiled function with - `np.random.default_rng(seed)` and pass it in. + `np.random.RandomState`. A jit-compiled caller passing anything + else gets a compile-time `numba.TypingError`. The pure-Python path + does no input checking -- it is a thin fallback, and its behaviour + for anything other than None or a `Generator` is unspecified and + unsupported. Build a generator outside any jit-compiled function + with `np.random.default_rng(seed)` and pass it in. If `rng` is None the draws come from a global random state that depends on the call site: NumPy's legacy global random state from @@ -237,8 +239,8 @@ def draw(cdf, size=None, rng=None): >>> rng = np.random.default_rng(1234) >>> qe.random.draw(cdf, 10, rng=rng) array([1, 0, 1, 0, 0, 0, 0, 0, 1, 0]) - >>> int(qe.random.draw(cdf, rng=np.random.default_rng(1234))) - 1 + >>> qe.random.draw(cdf, rng=np.random.default_rng(1234)) + np.int64(1) A `Generator` gives the same draws from Python and from within a jit-compiled function: @@ -255,12 +257,6 @@ def draw(cdf, size=None, rng=None): """ if rng is None: rng = np.random - elif not isinstance(rng, np.random.Generator): - raise TypeError( - 'quantecon.random.draw: `rng` must be None or an ' - f'np.random.Generator; got {rng!r}. Build a generator with ' - 'np.random.default_rng(seed) and pass that in.' - ) if isinstance(size, int): rs = rng.random(size) out = np.searchsorted(cdf, rs, side='right') From 4d83d40ed5c32947b9947fe32ce876e253a62aed Mon Sep 17 00:00:00 2001 From: Matt McKay Date: Sun, 16 Aug 2026 08:23:09 +1000 Subject: [PATCH 4/4] DOC: Drop Python-path behaviour from draw's Notes The Notes described the pure-Python fallback three times: an identical-behaviour promise across both call paths, an explicit "unspecified and unsupported" note about input checking, and the Python half of the rng=None global-state paragraph. Per review, say nothing about the Python version instead -- it is a thin fallback for when the JIT compiler is disabled, and it promises nothing. The Examples block asserting that Python and jitted give the same draws made the same promise in executable form, so it is now a plain jit-compiled usage example. The overload still must match the Python body, and TestDraw.test_python_jitted_agree still enforces that. Co-Authored-By: Claude Opus 5 (1M context) --- quantecon/random/utilities.py | 37 ++++++++++++----------------------- 1 file changed, 13 insertions(+), 24 deletions(-) diff --git a/quantecon/random/utilities.py b/quantecon/random/utilities.py index 1906369b..9ca73ece 100644 --- a/quantecon/random/utilities.py +++ b/quantecon/random/utilities.py @@ -190,7 +190,7 @@ def draw(cdf, size=None, rng=None): Random number generator to draw from. Must be a `np.random.Generator` or None; in particular, integer seeds and `np.random.RandomState` are not accepted. If None, the global - random state described in Notes is used. + random state is used (see Notes). Returns ------- @@ -199,29 +199,21 @@ def draw(cdf, size=None, rng=None): Notes ----- `draw` is intended primarily for use inside jit-compiled functions. - A `np.random.Generator` passed as `rng` behaves identically on both - paths: called from Python and called from within a jit-compiled - function it consumes the same generator and returns the same draws, - and the generator's state advances in place, so a single generator - can be shared between Python and jitted call sites and stays in - sync. Pass one whenever the draws should be reproducible. + Pass a `np.random.Generator` as `rng` whenever the draws should be + reproducible; the generator is consumed in place, so its state + advances across successive calls. `rng` accepts only None or a `Generator` because Numba's nopython mode can neither construct a generator from a seed nor represent `np.random.RandomState`. A jit-compiled caller passing anything - else gets a compile-time `numba.TypingError`. The pure-Python path - does no input checking -- it is a thin fallback, and its behaviour - for anything other than None or a `Generator` is unspecified and - unsupported. Build a generator outside any jit-compiled function - with `np.random.default_rng(seed)` and pass it in. - - If `rng` is None the draws come from a global random state that - depends on the call site: NumPy's legacy global random state from - Python, seeded by `np.random.seed`, and Numba's own internal random - state from within a jit-compiled function, seeded by calling - `np.random.seed` inside the jit-compiled function. Both are Mersenne - Twister and give the same stream for the same seed, but they advance - independently, so seeding one has no effect on the other. + else gets a compile-time `numba.TypingError`. Build a generator + outside any jit-compiled function with `np.random.default_rng(seed)` + and pass it in. + + If `rng` is None, a jit-compiled caller draws from Numba's own + internal random state, which is seeded by calling `np.random.seed` + inside a jit-compiled function; seeding NumPy's global random state + from Python has no effect on it. A single `Generator` must not be shared across the iterations of a `numba.prange` loop. Its state is updated in place, so concurrent @@ -242,15 +234,12 @@ def draw(cdf, size=None, rng=None): >>> qe.random.draw(cdf, rng=np.random.default_rng(1234)) np.int64(1) - A `Generator` gives the same draws from Python and from within a - jit-compiled function: + Inside a jit-compiled function: >>> from numba import njit >>> @njit ... def draw_jitted(cdf, size, rng): ... return qe.random.draw(cdf, size, rng) - >>> qe.random.draw(cdf, 5, np.random.default_rng(1234)) - array([1, 0, 1, 0, 0]) >>> draw_jitted(cdf, 5, np.random.default_rng(1234)) array([1, 0, 1, 0, 0])