update jax random key (#965)

kp992 · Copilot · mmcky · web-flow · commit 96ede29937a7 · 2026-08-12T16:55:22.000+10:00
Co-authored-by: copilot-swe-agent[bot] &lt;198982749+Copilot@users.noreply.github.com&gt;
Co-authored-by: mmcky &lt;8263752+mmcky@users.noreply.github.com&gt;
diff --git a/lectures/aiyagari_egm.md b/lectures/aiyagari_egm.md
@@ -472,7 +472,7 @@ Let's test it:
 ```{code-cell} ipython3
 household = create_household()
 prices = Prices(r=0.01, w=1.0)
-key = jax.random.PRNGKey(42)
+key = jax.random.key(42)
 
 with qe.Timer():
     K_supply = capital_supply(household, prices, key)
@@ -529,7 +529,7 @@ Let's compute the equilibrium:
 ```{code-cell} ipython3
 firm = Firm()
 household = create_household()
-key = jax.random.PRNGKey(42)
+key = jax.random.key(42)
 
 with qe.Timer():
     K_star = compute_equilibrium(firm, household, key)
diff --git a/lectures/ar1_bayes.md b/lectures/ar1_bayes.md
@@ -483,7 +483,7 @@ NUTS_kernel = numpyro.infer.NUTS(AR1_model)
 
 # Run MCMC
 mcmc = numpyro.infer.MCMC(NUTS_kernel, num_samples=50000, num_warmup=10000, progress_bar=False)
-mcmc.run(rng_key=random.PRNGKey(1), data=y)
+mcmc.run(rng_key=random.key(1), data=y)
 ```
 
 We plot the trace and posterior.
@@ -535,7 +535,7 @@ NUTS_kernel = numpyro.infer.NUTS(AR1_model_y0)
 
 # Run MCMC
 mcmc2 = numpyro.infer.MCMC(NUTS_kernel, num_samples=50000, num_warmup=10000, progress_bar=False)
-mcmc2.run(rng_key=random.PRNGKey(1), data=y)
+mcmc2.run(rng_key=random.key(1), data=y)
 ```
 
 Again we plot the trace and posterior.
diff --git a/lectures/back_prop.md b/lectures/back_prop.md
@@ -393,7 +393,7 @@ N = 3  # Number of layers
 layer_sizes = [1, ] * (N + 1)
 param_scale = 0.1
 step_size = 0.01
-params = init_network_params(layer_sizes, random.PRNGKey(1))
+params = init_network_params(layer_sizes, random.key(1))
 ```
 
 ```{code-cell} ipython3
@@ -496,12 +496,12 @@ f_val = f(grid)
 
 ```{code-cell} ipython3
 indices = jnp.arange(M)
-key = random.PRNGKey(0)
+key = random.key(0)
 
 def train(params, grid, f_val, key, num_epochs=300):
     for epoch in range(num_epochs):
         key, _ = random.split(key)
-        random_permutation = random.permutation(random.PRNGKey(1), indices)
+        random_permutation = random.permutation(random.key(1), indices)
         for x, y in zip(grid[random_permutation], f_val[random_permutation]):
             params = update_la(params, x, y)
             
diff --git a/lectures/ifp_advanced.md b/lectures/ifp_advanced.md
@@ -356,7 +356,7 @@ def create_ifp(
     assert β * ER < 1, "Stability condition failed."
 
     # Generate random draws using JAX
-    key = jax.random.PRNGKey(seed)
+    key = jax.random.key(seed)
     subkey1, subkey2 = jax.random.split(key)
     η_draws = jax.random.normal(subkey1, (shock_draw_size,))
     ζ_draws = jax.random.normal(subkey2, (shock_draw_size,))
@@ -586,7 +586,7 @@ def compute_asset_stationary(
     z_idx_0_vector = jnp.zeros(num_households).astype(jnp.int32)
 
     # Vectorize over many households
-    key = jax.random.PRNGKey(seed)
+    key = jax.random.key(seed)
     keys = jax.random.split(key, num_households)
     # Vectorize simulate_household in (key, a_0, z_idx_0)
     sim_all_households = jax.vmap(
diff --git a/lectures/ifp_egm.md b/lectures/ifp_egm.md
@@ -890,7 +890,7 @@ def compute_asset_stationary(
     z_idx_0_vector = jnp.zeros(num_households).astype(jnp.int32)
 
     # Vectorize over many households
-    key = jax.random.PRNGKey(seed)
+    key = jax.random.key(seed)
     keys = jax.random.split(key, num_households)
     # Vectorize simulate_household in (key, a_0, z_idx_0)
     sim_all_households = jax.vmap(
diff --git a/lectures/ifp_egm_transient_shocks.md b/lectures/ifp_egm_transient_shocks.md
@@ -411,7 +411,7 @@ def create_ifp(r=0.01,
                shock_draw_size=100,
                seed=1234):
 
-    key = jax.random.PRNGKey(seed)
+    key = jax.random.key(seed)
     s = jnp.linspace(0, savings_grid_max, savings_grid_size)
     Π, z_grid = jnp.array(Π), jnp.array(z_grid)
     R = 1 + r
@@ -779,7 +779,7 @@ def compute_asset_stationary(
     z_idx_0_vector = jnp.zeros(num_households).astype(jnp.int32)
 
     # Vectorize over many households
-    key = jax.random.PRNGKey(seed)
+    key = jax.random.key(seed)
     keys = jax.random.split(key, num_households)
     # Vectorize simulate_household in (key, a_0, z_idx_0)
     sim_all_households = jax.vmap(
diff --git a/lectures/lake_model.md b/lectures/lake_model.md
@@ -592,7 +592,7 @@ P = jnp.array([[1 - λ,        λ],
 xbar = rate_steady_state(model_markov)
 
 # Simulate the Markov chain - we need a different approach for random updates
-key = jax.random.PRNGKey(0)
+key = jax.random.key(0)
 
 def simulate_markov(P, initial_state, T, key):
     """Simulate Markov chain for T periods"""
diff --git a/lectures/mccall_fitted_vfi.md b/lectures/mccall_fitted_vfi.md
@@ -282,7 +282,7 @@ def create_mccall_model(
     ):
     """Factory function to create a McCall model instance."""
 
-    key = jax.random.PRNGKey(seed)
+    key = jax.random.key(seed)
     z_draws = jax.random.normal(key, (mc_size,))
 
     # Discretize just to get a suitable wage grid for interpolation
@@ -529,7 +529,7 @@ def simulate_employment_path(
     Simulate employment path for T periods starting from unemployment.
 
     """
-    key = jax.random.PRNGKey(seed)
+    key = jax.random.key(seed)
     c, α, β, ρ, ν, γ, w_grid, z_draws = model
 
     # Initial conditions: start unemployed with initial wage draw
@@ -678,7 +678,7 @@ def simulate_cross_section(
     """
     c, α, β, ρ, ν, γ, w_grid, z_draws = model
 
-    key = jax.random.PRNGKey(seed)
+    key = jax.random.key(seed)
 
     # Solve for optimal reservation wage
     w_bar = get_reservation_wage(model)
@@ -750,7 +750,7 @@ def plot_cross_sectional_unemployment(
     c, α, β, ρ, ν, γ, w_grid, z_draws = model
 
     # Get final employment state directly
-    key = jax.random.PRNGKey(42)
+    key = jax.random.key(42)
     w_bar = get_reservation_wage(model)
 
     # Initialize arrays
diff --git a/lectures/mccall_model.md b/lectures/mccall_model.md
@@ -802,7 +802,7 @@ class McCallModelContinuous(NamedTuple):
 def create_mccall_continuous(
         c=25, β=0.99, σ=0.5, μ=2.5, mc_size=1000, seed=1234
     ):
-    key = jax.random.PRNGKey(seed)
+    key = jax.random.key(seed)
     s = jax.random.normal(key, (mc_size,))
     w_draws = jnp.exp(μ + σ * s)
     return McCallModelContinuous(c, β, σ, μ, w_draws)
@@ -970,7 +970,7 @@ def simulate_lifetime_value(key, model, w_bar, n_periods=100):
 
     Parameters:
     -----------
-    key : jax.random.PRNGKey
+    key : jax.random.key
         Random key for JAX
     model : McCallModelContinuous
         The model containing parameters
@@ -1018,7 +1018,7 @@ def compute_mean_lifetime_value(model, w_bar, num_reps=10000, seed=1234):
     Compute mean lifetime value across many simulations.
 
     """
-    key = jax.random.PRNGKey(seed)
+    key = jax.random.key(seed)
     keys = jax.random.split(key, num_reps)
 
     # Vectorize the simulation across all replications
@@ -1096,7 +1096,7 @@ def compute_stopping_time_continuous(w_bar, key, model):
     -----------
     w_bar : float
         The reservation wage
-    key : jax.random.PRNGKey
+    key : jax.random.key
         Random key for JAX
     model : McCallModelContinuous
         The model containing wage draws
@@ -1148,7 +1148,7 @@ def compute_mean_stopping_time_continuous(w_bar, model, num_reps=100000, seed=12
         Average stopping time across all replications
     """
     # Generate a key for each MC replication
-    key = jax.random.PRNGKey(seed)
+    key = jax.random.key(seed)
     keys = jax.random.split(key, num_reps)
 
     # Vectorize compute_stopping_time_continuous and evaluate across keys
diff --git a/lectures/mccall_model_with_sep_markov.md b/lectures/mccall_model_with_sep_markov.md
@@ -605,7 +605,7 @@ def simulate_employment_path(
     Simulate employment path for T periods starting from unemployment.
 
     """
-    key = jax.random.PRNGKey(seed)
+    key = jax.random.key(seed)
     # Unpack model
     n, w_vals, P, P_cumsum, β, c, α, γ = model
 
@@ -813,7 +813,7 @@ def simulate_cross_section(
 
     Returns the cross-sectional unemployment rate.
     """
-    key = jax.random.PRNGKey(seed)
+    key = jax.random.key(seed)
 
     # Solve for optimal reservation wage
     v_u = vfi(model)
@@ -848,7 +848,7 @@ def plot_cross_sectional_unemployment(
 
     """
     # Get final employment state directly
-    key = jax.random.PRNGKey(42)
+    key = jax.random.key(42)
     v_u = vfi(model)
     w_bar = get_reservation_wage(v_u, model)
 
diff --git a/lectures/mccall_persist_trans.md b/lectures/mccall_persist_trans.md
@@ -192,7 +192,7 @@ class Model(NamedTuple):
     e_draws: jnp.ndarray
 
 def create_job_search_model(μ=0.0, s=1.0, d=0.0, ρ=0.9, σ=0.1, β=0.98, c=5.0, 
-                           mc_size=1000, grid_size=100, key=jax.random.PRNGKey(1234)):
+                           mc_size=1000, grid_size=100, key=jax.random.key(1234)):
     """
     Create a Model with computed grid and draws.
     """
@@ -377,7 +377,7 @@ def draw_duration(key, μ, s, d, ρ, σ, β, z_grid, f_star, t_max=10_000):
 
 
 def compute_unemployment_duration(
-        model, key=jax.random.PRNGKey(1234), num_reps=100_000
+        model, key=jax.random.key(1234), num_reps=100_000
     ):
     """
     Compute expected unemployment duration.
diff --git a/lectures/newton_method.md b/lectures/newton_method.md
@@ -825,7 +825,7 @@ The excess demand function is essentially the same, but now the matrix $A$ is $3
 dim = 3000
 
 # Create JAX random key
-key = jax.random.PRNGKey(0)
+key = jax.random.key(0)
 
 # Create a random matrix A and normalize the columns to sum to one
 A = jax.random.uniform(key, (dim, dim))
diff --git a/lectures/os_egm_jax.md b/lectures/os_egm_jax.md
@@ -117,7 +117,7 @@ def create_model(
     s_grid = jnp.linspace(1e-4, grid_max, grid_size)
 
     # Store shocks (with a seed, so results are reproducible)
-    key = jax.random.PRNGKey(seed)
+    key = jax.random.key(seed)
     shocks = jnp.exp(μ + s * jax.random.normal(key, shape=(shock_size,)))
 
     return Model(β, μ, s, s_grid, shocks, α)
diff --git a/lectures/unemployment_linear.md b/lectures/unemployment_linear.md
@@ -197,7 +197,7 @@ def run_nuts(model, data, seed=0, num_warmup=1000, num_samples=2000, num_chains=
                 num_warmup=num_warmup, num_samples=num_samples,
                 num_chains=num_chains, chain_method="vectorized",
                 progress_bar=False)
-    mcmc.run(random.PRNGKey(seed), jnp.asarray(data))
+    mcmc.run(random.key(seed), jnp.asarray(data))
     return mcmc
 ```
 
diff --git a/lectures/unemployment_shocks.md b/lectures/unemployment_shocks.md
@@ -204,7 +204,7 @@ def run_nuts(model, data, seed=0, num_warmup=2000, num_samples=4000, num_chains=
                 num_warmup=num_warmup, num_samples=num_samples,
                 num_chains=num_chains, chain_method="vectorized",
                 progress_bar=False)
-    mcmc.run(random.PRNGKey(seed), jnp.asarray(data))
+    mcmc.run(random.key(seed), jnp.asarray(data))
     return mcmc
 ```
 
diff --git a/scripts/test-jax-install.py b/scripts/test-jax-install.py
@@ -9,7 +9,7 @@ def matrix_multiply(a, b):
     return jnp.dot(a, b)
 
 # Example usage:
-key = jax.random.PRNGKey(0)
+key = jax.random.key(0)
 x = jax.random.normal(key, (1000, 1000))
 y = jax.random.normal(key, (1000, 1000))
 z = matrix_multiply(x, y)
