ENH: Extend random.draw to accept a random number generator

[mmcky]

Update (2026-08-15, edcd4c2): the argument is named rng and the documented contract is None | np.random.Generator, but it is enforced only at the jitted path (compile-time TypingError) — the pure-Python body is a thin unchecked fallback whose behaviour for other inputs is unspecified, per review. The requires-python bump is reverted (deferred to #869) and the numba floor is >=0.56.0, the minimum this feature needs. The decision table below records the original draft rationale; commit 59cf867 (2026-08-14) introduced the rng naming.

Closes #916. Opened as a draft because the issue asked for four design points to be settled before implementing — they are answered below and are the main thing worth reviewing.

quantecon.random.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 where it was called from, and a np.random.seed at Python level had no effect on the jitted path. This adds random_state, bringing draw into line with probvec, sample_without_replacement and DiscreteRV.draw.

Decisions on the four points raised in #916

Question	Decision	Reasoning
Parameter name: random_state or rng? Module-wide?	random_state, trailing: draw(cdf, size=None, random_state=None). Not a module-wide rename.	probvec, sample_without_replacement and DiscreteRV.draw already take random_state; draw was the one hole. Moving the whole module to an rng name is a contract change (it would drop RandomState support) and deserves its own issue and deprecation cycle rather than a drive-by on one function.
How do the pure-Python and @overload paths stay in agreement?	By test, not by construction. The jitted sized branch keeps its hand-written searchsorted loop rather than mirroring the Python body's vectorised call, and TestDraw.test_python_jitted_agree pins the two paths together across 3 seeds × 4 sizes.	Numba does support the array form, and the two produce bit-identical output — but the loop measured roughly 2× faster at size >= 1000 on this machine, so making the bodies textually identical would ship a silent performance regression under an "add random_state" release note.
Is check_random_state reachable from nopython mode?	No — the jitted path takes a narrower contract: None or a live np.random.Generator only. Everything else raises numba.TypingError at compile time, with a message that says what to pass instead.	check_random_state is pure Python, Numba cannot type np.random.RandomState at all, and it cannot construct a Generator from an integer seed inside nopython mode. There is no way to widen the jitted contract without silently remapping a seed to a different stream.
Backward compatibility	Purely additive; no deprecation.	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 old np.random.random(...), not merely statistically equivalent. The jitted None path still uses Numba's internal state, so a jitted caller seeding with np.random.seed inside the jitted function is unaffected. The new parameter is trailing-with-default and there was no third positional slot, so no existing call form changes meaning. All four pre-existing draw tests and both existing jitted wrappers pass unmodified.

A Generator passed into nopython mode reproduces the host stream bit-for-bit and its state is mutated in place, so one generator can be shared across Python and jitted calls and it stays in sync. That is now the recommended way to get reproducible draws from a call site that may be jit-compiled, and it is what the new docstring example demonstrates.

Points I would like a decision on

The numba>=0.49.0 → numba>=0.59.0 floor bump in pyproject.toml. This change needs types.NumPyRandomGeneratorType, which numba first shipped in 0.56.0 (absent in 0.55.0). I set the floor at 0.59 rather than 0.56 because 0.59 is the first release with cp312 wheels, and the classifiers and the CI matrix already declare 3.12/3.13/3.14 — a floor below that names versions installable on no Python this package supports. Happy to drop this hunk if you would rather handle dependency metadata separately; the alternative of a getattr guard is worse, since it would silently degrade the feature into "every random_state is rejected" on an old numba rather than failing loudly. Note requires-python = ">=3.7" two lines above is stale by the same argument, but #864 already changes that line — see below.

Merge order against #864. #864 (docs/consistency-audit) rewrites this same docstring: it added a Notes block in 08682a7 documenting the two-generator behaviour, which #916 itself says should be replaced once this lands. That block is superseded here. Both PRs also touch pyproject.toml within three lines of each other. Whoever merges second gets a small, obvious conflict in quantecon/random/utilities.py and possibly pyproject.toml. My suggestion is to land #864 first and rebase this on top, since #864 is the larger diff. I left the one-line quantecon/random/__init__.py docstring fix (3. draw) out of this PR because #864 already makes exactly that change.

Out of scope, but found along the way

The pure-Python body dispatches on isinstance(size, int) while the overload dispatches on types.Integer. These disagree for numpy integers: qe.random.draw(cdf, np.int64(10)) silently returns one scalar from Python and a 10-element array from jitted code, with no error. size=True is the mirror image (Python raises TypeError, jitted returns a scalar). Both predate this change and are untouched here — fixing them is a silent behaviour change that would muddy this PR's release note and its bisect story. Tracked separately in #918.

Tests

Thirteen new tests in quantecon/random/tests/test_utilities.py, covering: Python/jitted agreement for the same Generator; reproducibility on both paths; in-place generator state mutation across a mixed chain of jitted and Python calls; that exactly one variate is consumed per draw, so a shared generator does not desynchronise; the keyword call form from jitted code; int-seed and RandomState on the Python path; that a seed and a default_rng of the same integer differ; byte-identity of the None path against the legacy global stream; that all three Numba spellings of "no random_state" compile (omitted arrives as Python None, explicit None and a forwarded default both arrive as types.none); that in-jit np.random.seed is still reproducible; and the three compile-time rejections.

The rejection tests assert on tokens that can only come from draw's own message. A plain 'random_state' in str(e) looks right but is vacuous — Numba echoes the caller's source line alongside the error, so the parameter name appears in the message whether or not the implementation produced it.

Each new test was checked by mutation: reverting the Generator branch to np.random, removing check_random_state, dropping either clause of the None predicate, swapping the two overload branches, and over-consuming a variate on any of the four Generator branches each fail at least one test.

Verification

Run on numba 0.62.1 / numpy 2.3.5 / python 3.13.9:

• pytest quantecon — 613 passed
• pytest quantecon/random/tests/test_utilities.py — 22 passed
• NUMBA_DISABLE_JIT=1 pytest quantecon/random/tests/test_utilities.py — 19 passed, 3 skipped (the three compile-time rejection tests are skipif-guarded, since with JIT off the @njit wrappers run the Python body and check_random_state accepts the int seed)
• flake8 --select=F401,F405,E231 quantecon — clean
• Full flake8 on both touched files — clean apart from the pre-existing E305 at utilities.py:92
• The new draw doctest passes under --doctest-modules. Note doctests are not run in CI, and the two sibling examples in this module fail there today for a pre-existing reason (they assume an ambient qe); the new example is self-contained.

🤖 Generated with Claude Code

[coveralls]

coverage: 90.605% (+0.04%) from 90.57% — enh/draw-random-state-916 into main

[mmcky]

@oyamad just started work on this. This will remain in draft until I get a chance to do a detailed review. I'm not sure we need all these different states supported and we may just upgrade / deprecate.

[Copilot]

🟡 Not ready to approve

The dependency floor bump introduces packaging metadata inconsistency (requires-python vs classifiers/numba floor) and there are a couple of user-facing docstring text issues to fix.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

This review doesn't count toward merge requirements. Sign up for the private preview to control whether Copilot approvals count.

Pull request overview

Extends quantecon.random.draw to accept an explicit RNG via a new trailing random_state parameter, aligning its behavior (and reproducibility guarantees) with other quantecon.random APIs and ensuring consistent results between pure-Python and Numba-jitted call sites when a np.random.Generator is provided.

Changes:

• Add random_state to quantecon.random.draw, using check_random_state on the Python path and supporting np.random.Generator (or None) in the Numba overload.
• Add comprehensive tests validating Python/jitted agreement, generator state advancement, and compile-time rejections in nopython mode.
• Bump the numba minimum version in package dependencies to support types.NumPyRandomGeneratorType.

File summaries

File	Description
quantecon/random/utilities.py	Adds random_state to draw, updates docstring, and extends the Numba overload to accept np.random.Generator in nopython mode.
quantecon/random/tests/test_utilities.py	Adds new test coverage for RNG behavior and Python/jitted agreement for draw.
pyproject.toml	Raises the minimum numba dependency version.

Review details

• Files reviewed: 3/3 changed files
• Comments generated: 3
• Review effort level: Lite

We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.

[oyamad]

@mmcky Thanks!

Here's my argument:

1. I would like to take this stance: draw is primarily for use in a jitted function; it "happens" to work in pure Python, but how it behaves is implementation detail, and thus unspecified/unsupported by the public API.
2. We are now to encourage the use pattern, rng = np.random.default_rng(1234); function_with_randomness(..., rng=rng). Thus add a new optional argument rng, where we don't need to (or shouldn't) try to mimic the existing random_state, so it only accepts None (default) or Generator; in particular, don't accept integer. This will simplify some of the issues.
3. [This was pointed out by Claude Fable 5] The policy above is stricter than "SPEC 7", which recommends normalizing rng via np.random.default_rng(rng) (hence accepting int seeds). That normalization step is not possible in nopython mode. So my proposal is: in Numba functions, take rng as None | Generator only.
4. Regarding merge order: The only restriction is that doctest be activated after this PR is merged.
5. On random.draw: size dispatch diverges between the Python and jitted paths for numpy integers #918: That is a bug I introduced when I rewrote the code with @overload... Just replace if isinstance(size, int): with if isinstance(size, (int, np.integer)):. That's it.

[mmcky]

@oyamad Thanks — all five points adopted in 59cf867.

Points 1–3 (rng, None | Generator only). The argument is renamed to rng and the narrow contract now holds on both paths, not just the jitted one: the Python body raises TypeError for anything that isn't None or a np.random.Generator, with the same "build one with np.random.default_rng(seed)" guidance the jitted path gives in its compile-time TypingError. check_random_state is no longer involved. I took your point 3 as implying the Python path should be equally strict — accepting seeds on one path only would let the two paths disagree, which is the thing this PR exists to prevent. The docstring now leads with your point 1 (intended primarily for jitted use), and the int-seed/RandomState material is gone from Notes; the example uses a shared Generator.

Tests: the two seed/RandomState acceptance tests are gone, and the rejection tests now run on both paths (they no longer need skipif under NUMBA_DISABLE_JIT, since the Python body rejects too). Full suite: 611 passed; with JIT disabled: 19 passed, 1 skipped; the new doctest passes.

Point 4 (merge order). Understood — this PR first, doctest activation after. Since this now lands first, I also bumped requires-python to >=3.12 here (Copilot flagged the inconsistency with the new numba floor); the doctest branch gets a trivial rebase conflict on that line.

Point 5 (#918). I'll open a separate one-line PR with isinstance(size, (int, np.integer)) as you suggest, keeping it out of this diff.

One follow-up worth tracking: draw is now the one function in quantecon.random taking rng while probvec and sample_without_replacement take random_state. If the direction is a SPEC-7-style module-wide move to rng, I'm happy to open an issue proposing that migration (with a deprecation cycle for random_state) so the naming split is deliberate rather than accidental.

[oyamad]

@mmcky Thanks for the revision.

My earlier suggestion was: leave the behavior of the Python version unspecified/unsupported, by which I meant "we promise nothing". So I wouldn't do any input check, and the following would be enough:

def draw(cdf, size=None, rng=None):
    if rng is None:
        rng = np.random
    r = rng.random()
    return np.searchsorted(cdf, r, side='right')

In particular, we don't need to reject RandomState. This way we can keep the Python version as thin as possible. What do you think?

[mmcky]

Agreed — done in edcd4c2. The Python body is now the thin version (no input check; the size branch stays). The None | Generator contract is enforced only by the overload's TypingError, and the docstring now says the Python path is an unchecked fallback, unspecified for other inputs. The two Python-path rejection tests are removed and the three compile-time ones are again skipped under NUMBA_DISABLE_JIT. Full suite: 611 passed; with JIT off: 17 passed, 3 skipped.

[oyamad]

@mmcky Thanks! This is what I had in mind.

One minor comment: In the Notes, I would not mention the behavior of the Python version at all. But not blocking.

[mmcky]

Done in 4d83d40 — the Notes no longer describe the Python path at all. I also dropped the Python-vs-jitted line from the Examples, since it made the same promise; test_python_jitted_agree still holds the two paths together internally.

[mmcky]

thanks @oyamad for all of your valuable review comments.