QuantEcon / QuantEcon/lecture-python.myst

Modernize jax.random.PRNGKey to jax.random.key across lectures

Open
#914 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement
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.key is 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 raw uint32[2] representation of a key.
  • New lectures (unemployment_bayes #910, population_ssm #911) will use random.key from the start.

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.