Skip to content

ENH: Extend random.draw to accept a random number generator - #917

Draft
mmcky wants to merge 1 commit into
mainfrom
enh/draw-random-state-916
Draft

ENH: Extend random.draw to accept a random number generator#917
mmcky wants to merge 1 commit into
mainfrom
enh/draw-random-state-916

Conversation

@mmcky

@mmcky mmcky commented Aug 1, 2026

Copy link
Copy Markdown
Contributor

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.0numba>=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

`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) <noreply@anthropic.com>
@coveralls

Copy link
Copy Markdown

Coverage Status

coverage: 90.58% (+0.03%) from 90.546% — enh/draw-random-state-916 into main

@mmcky

mmcky commented Aug 1, 2026

Copy link
Copy Markdown
Contributor Author

@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 AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 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.

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.
Comment on lines 5 to 9
---------
probvec
sample_without_replacement
draw

Comment thread pyproject.toml
Comment on lines 24 to 27
requires-python = ">=3.7"
dependencies = [
'numba>=0.49.0',
'numba>=0.59.0',
'numpy>=1.17.0',
@oyamad

oyamad commented Aug 1, 2026

Copy link
Copy Markdown
Member

@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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Extend random.draw to accept a random number generator

4 participants