ENH: Extend random.draw to accept a random number generator - #917
Conversation
`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>
|
@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. |
There was a problem hiding this comment.
🟡 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_statetoquantecon.random.draw, usingcheck_random_stateon the Python path and supportingnp.random.Generator(orNone) in the Numba overload. - Add comprehensive tests validating Python/jitted agreement, generator state advancement, and compile-time rejections in nopython mode.
- Bump the
numbaminimum version in package dependencies to supporttypes.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. |
| --------- | ||
| probvec | ||
| sample_without_replacement | ||
| draw | ||
|
|
| requires-python = ">=3.7" | ||
| dependencies = [ | ||
| 'numba>=0.49.0', | ||
| 'numba>=0.59.0', | ||
| 'numpy>=1.17.0', |
|
@mmcky Thanks! Here's my argument:
|
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.drawwas the only function inquantecon.randomwith norandom_stateargument. It callednp.random.random()directly in both the pure-Python body and the@overloadimplementation, so the generator it used depended on where it was called from, and anp.random.seedat Python level had no effect on the jitted path. This addsrandom_state, bringingdrawinto line withprobvec,sample_without_replacementandDiscreteRV.draw.Decisions on the four points raised in #916
random_stateorrng? Module-wide?random_state, trailing:draw(cdf, size=None, random_state=None). Not a module-wide rename.probvec,sample_without_replacementandDiscreteRV.drawalready takerandom_state;drawwas the one hole. Moving the whole module to anrngname is a contract change (it would dropRandomStatesupport) and deserves its own issue and deprecation cycle rather than a drive-by on one function.@overloadpaths stay in agreement?searchsortedloop rather than mirroring the Python body's vectorised call, andTestDraw.test_python_jitted_agreepins the two paths together across 3 seeds × 4 sizes.size >= 1000on this machine, so making the bodies textually identical would ship a silent performance regression under an "addrandom_state" release note.check_random_statereachable from nopython mode?Noneor a livenp.random.Generatoronly. Everything else raisesnumba.TypingErrorat compile time, with a message that says what to pass instead.check_random_stateis pure Python, Numba cannot typenp.random.RandomStateat all, and it cannot construct aGeneratorfrom an integer seed inside nopython mode. There is no way to widen the jitted contract without silently remapping a seed to a different stream.check_random_state(None)returnsnp.random.mtrand._randby identity and is non-consuming, so the PythonNonepath is byte-identical to the oldnp.random.random(...), not merely statistically equivalent. The jittedNonepath still uses Numba's internal state, so a jitted caller seeding withnp.random.seedinside 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-existingdrawtests and both existing jitted wrappers pass unmodified.A
Generatorpassed 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.0floor bump inpyproject.toml. This change needstypes.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 agetattrguard is worse, since it would silently degrade the feature into "everyrandom_stateis rejected" on an old numba rather than failing loudly. Noterequires-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 touchpyproject.tomlwithin three lines of each other. Whoever merges second gets a small, obvious conflict inquantecon/random/utilities.pyand possiblypyproject.toml. My suggestion is to land #864 first and rebase this on top, since #864 is the larger diff. I left the one-linequantecon/random/__init__.pydocstring 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 ontypes.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=Trueis the mirror image (Python raisesTypeError, 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 andRandomStateon the Python path; that a seed and adefault_rngof the same integer differ; byte-identity of theNonepath against the legacy global stream; that all three Numba spellings of "no random_state" compile (omitted arrives as PythonNone, explicitNoneand a forwarded default both arrive astypes.none); that in-jitnp.random.seedis 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, removingcheck_random_state, dropping either clause of theNonepredicate, 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 passedpytest quantecon/random/tests/test_utilities.py— 22 passedNUMBA_DISABLE_JIT=1 pytest quantecon/random/tests/test_utilities.py— 19 passed, 3 skipped (the three compile-time rejection tests areskipif-guarded, since with JIT off the@njitwrappers run the Python body andcheck_random_stateaccepts the int seed)flake8 --select=F401,F405,E231 quantecon— cleanflake8on both touched files — clean apart from the pre-existingE305atutilities.py:92drawdoctest 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 ambientqe); the new example is self-contained.🤖 Generated with Claude Code