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

[jstac]

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 [unemployment_bayes] New lecture: Bayesian estimation of nonlinear unemployment dynamics #910, population_ssm [population_ssm] New lecture: Bayesian state-space models (linear filtering to nonlinear population dynamics) #911) will use random.key from the start.