chore: replace legacy jax.random.PRNGKey with modern jax.random.key (#2134)

* chore: replace legacy `jax.random.PRNGKey` with modern `jax.random.key`

* fix: files reformatted with `ruff==0.15.0`

* test: update PRNG key test

* test: update test effected by the `optax==0.2.7` release

* Revert "test: update test effected by the `optax==0.2.7` release"

This reverts commit 14d3ee8.

* Revert "fix: files reformatted with `ruff==0.15.0`"

This reverts commit 2482781.

* test: pickle `jax.random.key`

* test: use legacy keys in `test/test_pickle.py` until
jax-ml/jax#35065 is resolved

Author: Qazalbash

README.md

@@ -62,7 +62,7 @@ Let us infer the values of the unknown parameters in our model by running MCMC u
 
 >>> nuts_kernel = NUTS(eight_schools)
 >>> mcmc = MCMC(nuts_kernel, num_warmup=500, num_samples=1000)
->>> rng_key = random.PRNGKey(0)
+>>> rng_key = random.key(0)
 >>> mcmc.run(rng_key, J, sigma, y=y, extra_fields=('potential_energy',))
 
 ```
@@ -111,7 +111,7 @@ The values above 1 for the split Gelman Rubin diagnostic (`r_hat`) indicates tha
 
 >>> nuts_kernel = NUTS(eight_schools_noncentered)
 >>> mcmc = MCMC(nuts_kernel, num_warmup=500, num_samples=1000)
->>> rng_key = random.PRNGKey(0)
+>>> rng_key = random.key(0)
 >>> mcmc.run(rng_key, J, sigma, y=y, extra_fields=('potential_energy',))
 >>> mcmc.print_summary(exclude_deterministic=False)  # doctest: +SKIP
 
@@ -161,7 +161,7 @@ Now, let us assume that we have a new school for which we have not observed any
 ...     return numpyro.sample('obs', dist.Normal(mu, tau))
 
 >>> predictive = Predictive(new_school, mcmc.get_samples())
->>> samples_predictive = predictive(random.PRNGKey(1))
+>>> samples_predictive = predictive(random.key(1))
 >>> print(np.mean(samples_predictive['obs']))  # doctest: +SKIP
 3.9886456
 
@@ -286,18 +286,18 @@ conda install -c conda-forge numpyro
 
 1. Unlike in Pyro, `numpyro.sample('x', dist.Normal(0, 1))` does not work. Why?
 
-   You are most likely using a `numpyro.sample` statement outside an inference context. JAX does not have a global random state, and as such, distribution samplers need an explicit random number generator key ([PRNGKey](https://jax.readthedocs.io/en/latest/jax.random.html#jax.random.PRNGKey)) to generate samples from. NumPyro's inference algorithms use the [seed](https://num.pyro.ai/en/latest/handlers.html#seed) handler to thread in a random number generator key, behind the scenes.
+   You are most likely using a `numpyro.sample` statement outside an inference context. JAX does not have a global random state, and as such, distribution samplers need an explicit random number generator key ([PRNG Key](https://jax.readthedocs.io/en/latest/jax.random.html#jax.random.key)) to generate samples from. NumPyro's inference algorithms use the [seed](https://num.pyro.ai/en/latest/handlers.html#seed) handler to thread in a random number generator key, behind the scenes.
 
    Your options are:
 
-   - Call the distribution directly and provide a `PRNGKey`, e.g. `dist.Normal(0, 1).sample(PRNGKey(0))`
-   - Provide the `rng_key` argument to `numpyro.sample`. e.g. `numpyro.sample('x', dist.Normal(0, 1), rng_key=PRNGKey(0))`.
+   - Call the distribution directly and provide a PRNG key, e.g. `dist.Normal(0, 1).sample(key(0))`
+   - Provide the `rng_key` argument to `numpyro.sample`. e.g. `numpyro.sample('x', dist.Normal(0, 1), rng_key=key(0))`.
    - Wrap the code in a `seed` handler, used either as a context manager or as a function that wraps over the original callable. e.g.
 
         ```python
-        with handlers.seed(rng_seed=0):  # random.PRNGKey(0) is used
-            x = numpyro.sample('x', dist.Beta(1, 1))    # uses a PRNGKey split from random.PRNGKey(0)
-            y = numpyro.sample('y', dist.Bernoulli(x))  # uses different PRNGKey split from the last one
+        with handlers.seed(rng_seed=0):  # random.key(0) is used
+            x = numpyro.sample('x', dist.Beta(1, 1))    # uses a PRNG key split from random.key(0)
+            y = numpyro.sample('y', dist.Bernoulli(x))  # uses different PRNG key split from the last one
         ```
 
      , or as a higher order function:

docs/source/contrib.rst

@@ -183,7 +183,7 @@ Other kernels can be used similarly.
         ...     return x, y_obs
 
 
-        >>> rng_key = random.PRNGKey(seed=42)
+        >>> rng_key = random.key(seed=42)
         >>> rng_key, rng_subkey = random.split(rng_key)
         >>> x, y_obs = generate_synthetic_data(
         ...     rng_key=rng_subkey, start=0, stop=1, num=80, scale=0.3

examples/annotation.py

@@ -284,12 +284,12 @@ def main(args):
         num_chains=args.num_chains,
         progress_bar=False if "NUMPYRO_SPHINXBUILD" in os.environ else True,
     )
-    mcmc.run(random.PRNGKey(0), *data)
+    mcmc.run(random.key(0), *data)
     mcmc.print_summary()
 
     posterior_samples = mcmc.get_samples()
     predictive = Predictive(model, posterior_samples, infer_discrete=True)
-    discrete_samples = predictive(random.PRNGKey(1), *data)
+    discrete_samples = predictive(random.key(1), *data)
 
     item_class = vmap(lambda x: jnp.bincount(x, length=4), in_axes=1)(
         discrete_samples["c"].squeeze(-1)

examples/ar2.py

@@ -105,7 +105,7 @@ def run_inference(model, args, rng_key, y):
 def main(args):
     # generate artificial dataset
     num_data = args.num_data
-    rng_key = jax.random.PRNGKey(0)
+    rng_key = jax.random.key(0)
     t = jnp.arange(0, num_data)
     y = jnp.sin(t) + random.normal(rng_key, (num_data,)) * 0.1
 

examples/baseball.py

@@ -195,7 +195,7 @@ def main(args):
     for i, model in enumerate(
         (fully_pooled, not_pooled, partially_pooled, partially_pooled_with_logit)
     ):
-        rng_key, rng_key_predict = random.split(random.PRNGKey(i + 1))
+        rng_key, rng_key_predict = random.split(random.key(i + 1))
         zs = run_inference(model, at_bats, hits, rng_key, args)
         predict(model, at_bats, hits, zs, rng_key_predict, player_names)
         predict(

examples/bnn.py

@@ -126,7 +126,7 @@ def main(args):
     X, Y, X_test = get_data(N=N, D_X=D_X)
 
     # do inference
-    rng_key, rng_key_predict = random.split(random.PRNGKey(0))
+    rng_key, rng_key_predict = random.split(random.key(0))
     samples = run_inference(model, args, rng_key, X, Y, D_H)
 
     # predict Y_test at inputs X_test

examples/capture_recapture.py

@@ -343,7 +343,7 @@ def main(args):
     )
 
     model = models[args.model]
-    rng_key = random.PRNGKey(args.rng_seed)
+    rng_key = random.key(args.rng_seed)
     run_inference(model, capture_history, sex, rng_key, args)
 
 

examples/covtype.py

@@ -72,7 +72,7 @@ def model(data, labels, subsample_size=None):
 
 
 def benchmark_hmc(args, features, labels):
-    rng_key = random.PRNGKey(1)
+    rng_key = random.key(1)
     start = time.time()
     # a MAP estimate at the following source
     # https://github.com/google/edward2/blob/master/examples/no_u_turn_sampler/logistic_regression.py#L117
@@ -174,7 +174,7 @@ def benchmark_hmc(args, features, labels):
         subsample_size = 1000
         guide = AutoBNAFNormal(model, num_flows=1, hidden_factors=[8])
         svi = SVI(model, guide, numpyro.optim.Adam(0.01), Trace_ELBO())
-        svi_result = svi.run(random.PRNGKey(2), 2000, features, labels)
+        svi_result = svi.run(random.key(2), 2000, features, labels)
         params, losses = svi_result.params, svi_result.losses
         plt.plot(losses)
         plt.show()

examples/cvae-flax/train_baseline.py

@@ -10,7 +10,7 @@
 
 
 def create_train_state(model, x, learning_rate_fn):
-    params = model.init(random.PRNGKey(0), x)
+    params = model.init(random.key(0), x)
     tx = optax.adam(learning_rate_fn)
     state = TrainState.create(apply_fn=model.apply, params=params, tx=tx)
     return state
@@ -65,7 +65,7 @@ def train_baseline(
 ):
     state = create_train_state(model, train_fetch(0, train_idx)[0], 0.003)
 
-    rng = random.PRNGKey(0)
+    rng = random.key(0)
     best_val_loss = jnp.inf
     best_state = state
     for i in range(n_epochs):

examples/cvae-flax/train_cvae.py

@@ -78,14 +78,14 @@ def train_cvae(
     n_epochs=100,
 ):
     svi, state = create_train_state(
-        random.PRNGKey(23), model, guide, train_fetch, baseline_params, 0.003
+        random.key(23), model, guide, train_fetch, baseline_params, 0.003
     )
 
     p1 = baseline_params.unfreeze()["params"]["Dense_0"]["kernel"]
     p2 = state.optim_state[1][0]["baseline$params"]["Dense_0"]["kernel"]
     assert jnp.all(p1 == p2)
 
-    rng = random.PRNGKey(0)
+    rng = random.key(0)
     best_val_loss = jnp.inf
     best_state = state
     for i in range(n_epochs):