QuantEcon / QuantEcon/lecture-python.myst
Modernize jax.random.PRNGKey to jax.random.key across lectures
Nobody has claimed this yet.
- Dominant language
- TeX
- Stars
- 123
- Forks
- 57
- Avg merge
- 3d 10h
- Merged PRs (30d)
- 11
Description
Goal
Modernize jax.random.PRNGKey(...) to jax.random.key(...) across the lecture series.
Background
random.PRNGKey is the legacy key constructor (returns a raw uint32[2] array). JAX now recommends random.key, which returns a typed key (dtype=key<...>). PRNGKey is not formally deprecated yet (no warning as of JAX 0.8.0), so this is forward-looking cleanup rather than a fix — but doing it now avoids churn when JAX eventually deprecates the legacy API.
Verified while rewriting bayes_nonconj.md (#913) that NumPyro accepts the typed key throughout (MCMC, SVI, sample_posterior), and the lecture runs end-to-end with it. bayes_nonconj.md already uses random.key.
Scope
Files still using PRNGKey (grep -rl "PRNGKey" lectures/*.md):
-
lectures/ar1_bayes.md -
lectures/mix_model.md -
lectures/aiyagari_egm.md -
lectures/back_prop.md -
lectures/ifp_advanced.md -
lectures/ifp_egm.md -
lectures/ifp_egm_transient_shocks.md -
lectures/lake_model.md -
lectures/mccall_model.md -
lectures/mccall_model_with_sep_markov.md -
lectures/mccall_fitted_vfi.md -
lectures/mccall_persist_trans.md -
lectures/mccall_persist_trans.md -
lectures/os_egm_jax.md -
lectures/newton_method.md
Notes
- This spans both NumPyro lectures (
ar1_bayes,mix_model) and plain-JAX lectures (the rest);random.keyis the right call in both cases. - Mechanical change (
random.PRNGKey(->random.key(), plus any prose references), but each lecture should be re-executed to confirm nothing downstream relies on the rawuint32[2]representation of a key. - New lectures (
unemployment_bayes#910,population_ssm#911) will userandom.keyfrom the start.
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with the listed lectures and use the provided grep for remaining PRNGKey references. Replace the legacy constructor and related prose, then re-execute each affected lecture, including the NumPyro examples, to confirm typed keys work throughout. Done means all listed files are updated and their executions complete successfully.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- documentation
- Issue type
- Refactor
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 68/100