Update API usage

README.md

@@ -77,7 +77,7 @@ initial_position = {"loc": 1., "scale": 2.}
 state = nuts.init(initial_position)
 
 # Iterate
-rng_key = jax.random.PRNGKey(0)
+rng_key = jax.random.key(0)
 for _ in range(100):
     rng_key, nuts_key = jax.random.split(rng_key)
     state, _ = nuts.step(nuts_key, state)

docs/examples/howto_custom_gradients.md

@@ -184,7 +184,7 @@ def logdensity_fn(y):
 hmc = blackjax.hmc(logdensity_fn,1e-3, jnp.ones(1), 10)
 state = hmc.init(1.)
 
-rng_key = jax.random.PRNGKey(0)
+rng_key = jax.random.key(0)
 new_state, info = hmc.step(rng_key, state)
 ```
 

docs/examples/howto_metropolis_within_gibbs.md

@@ -200,7 +200,7 @@ def sampling_loop(rng_key, initial_state, parameters, num_samples):
 
 ```{code-cell} ipython3
 %%time
-rng_key = jax.random.PRNGKey(0)
+rng_key = jax.random.key(0)
 positions = sampling_loop(rng_key, initial_state, parameters, 10_000)
 ```
 
@@ -305,7 +305,7 @@ def sampling_loop_general(rng_key, initial_state, logdensity_fn, step_fn, init,
 
 ```{code-cell} ipython3
 %%time
-rng_key = jax.random.PRNGKey(0)
+rng_key = jax.random.key(0)
 positions_general = sampling_loop_general(
     rng_key=rng_key,
     initial_state=initial_state,

docs/examples/howto_other_frameworks.md

@@ -138,7 +138,7 @@ step_size=1e-3
 nuts = blackjax.nuts(numba_logpdf, step_size, inverse_mass_matrix)
 init = nuts.init(0.)
 
-rng_key = jax.random.PRNGKey(0)
+rng_key = jax.random.key(0)
 state, info = nuts.step(rng_key, init)
 
 for _ in range(10):

docs/examples/howto_sample_multiple_chains.md

@@ -96,7 +96,7 @@ And finally, to put `jax.vmap` and `jax.pmap` on an equal foot we sample as many
 import multiprocessing
 
 
-rng_key = jax.random.PRNGKey(0)
+rng_key = jax.random.key(0)
 num_chains = multiprocessing.cpu_count()
 ```
 

docs/examples/howto_use_aesara.md

@@ -136,7 +136,7 @@ def init_param_fn(seed):
         "thetas": jax.random.uniform(seed, (n_rat_tumors,), "float64", minval=0, maxval=1),
     }
 
-rng_key = jax.random.PRNGKey(0)
+rng_key = jax.random.key(0)
 init_position = init_param_fn(rng_key)
 ```
 

docs/examples/howto_use_numpyro.md

@@ -67,7 +67,7 @@ import jax
 
 from numpyro.infer.util import initialize_model
 
-rng_key = jax.random.PRNGKey(0)
+rng_key = jax.random.key(0)
 init_params, potential_fn_gen, *_ = initialize_model(
     rng_key,
     eight_schools_noncentered,

docs/examples/howto_use_oryx.md

@@ -109,7 +109,7 @@ from oryx.core.ppl import joint_sample
 
 
 bnn = mlp([50, 50], num_classes)
-initial_weights = joint_sample(bnn)(jax.random.PRNGKey(0), jnp.ones(num_features))
+initial_weights = joint_sample(bnn)(jax.random.key(0), jnp.ones(num_features))
 
 print(initial_weights.keys())
 ```
@@ -136,7 +136,7 @@ We can now run the window adaptation to get good values for the parameters of th
 %%time
 import blackjax
 
-rng_key = jax.random.PRNGKey(0)
+rng_key = jax.random.key(0)
 adapt = blackjax.window_adaptation(blackjax.nuts, logdensity_fn)
 (last_state, parameters), _ = adapt.run(rng_key, initial_weights, 100)
 kernel = blackjax.nuts(logdensity_fn, **parameters).step
@@ -173,7 +173,7 @@ posterior_weights = states.position
 
 output_logits = jax.vmap(
     lambda weights: jax.vmap(lambda x: intervene(bnn, **weights)(
-        jax.random.PRNGKey(0), x)
+        jax.random.key(0), x)
     )(features)
 )(posterior_weights)
 

docs/examples/howto_use_pymc.md

@@ -67,7 +67,7 @@ import jax
 init_position_dict = model.initial_point()
 init_position = [init_position_dict[rv] for rv in rvs]
 
-rng_key = jax.random.PRNGKey(1234)
+rng_key = jax.random.key(1234)
 
 adapt = blackjax.window_adaptation(blackjax.nuts, logdensity_fn)
 (last_state, parameters), _ = adapt.run(rng_key, init_position, 1000)

docs/examples/howto_use_tfp.md

@@ -87,7 +87,7 @@ initial_position = {
 }
 
 
-rng_key = jax.random.PRNGKey(0)
+rng_key = jax.random.key(0)
 adapt = blackjax.window_adaptation(
     blackjax.hmc, logdensity_fn, num_integration_steps=3
 )

docs/examples/quickstart.md

@@ -97,7 +97,7 @@ def inference_loop(rng_key, kernel, initial_state, num_samples):
 
 ```{code-cell} python
 %%time
-rng_key = jax.random.PRNGKey(0)
+rng_key = jax.random.key(0)
 states = inference_loop(rng_key, hmc_kernel, initial_state, 10_000)
 
 loc_samples = states.position["loc"].block_until_ready()
@@ -136,7 +136,7 @@ initial_state
 
 ```{code-cell} python
 %%time
-rng_key = jax.random.PRNGKey(0)
+rng_key = jax.random.key(0)
 states = inference_loop(rng_key, nuts.step, initial_state, 4_000)
 
 loc_samples = states.position["loc"].block_until_ready()

docs/index.md

@@ -37,7 +37,7 @@ initial_position = {"loc": 1., "scale": 2.}
 state = nuts.init(initial_position)
 
 # Iterate
-rng_key = jax.random.PRNGKey(0)
+rng_key = jax.random.key(0)
 step = jax.jit(nuts.step)
 for _ in range(1_000):
    rng_key, nuts_key = jax.random.split(rng_key)

tests/adaptation/test_step_size.py

@@ -16,7 +16,7 @@ def test_reasonable_step_size(self):
         def logdensity_fn(x):
             return -jnp.sum(0.5 * x)
 
-        rng_key = jax.random.PRNGKey(0)
+        rng_key = jax.random.key(0)
         run_key0, run_key1 = jax.random.split(rng_key, 2)
 
         init_position = jnp.array([3.0])

tests/mcmc/test_latent_gaussian.py

@@ -14,7 +14,7 @@ class GaussianTest(chex.TestCase):
     def test_gaussian(self, seed, mean):
         n_samples = 500_000
 
-        key = jax.random.PRNGKey(seed)
+        key = jax.random.key(seed)
         key1, key2, key3, key4, key5 = jax.random.split(key, 5)
 
         D = 5

tests/mcmc/test_metrics.py

@@ -11,7 +11,7 @@
 class GaussianEuclideanMetricsTest(chex.TestCase):
     def setUp(self):
         super().setUp()
-        self.key = random.PRNGKey(0)
+        self.key = random.key(0)
         self.dtype = "float32"
 
     @parameterized.named_parameters(

tests/mcmc/test_proposal.py

@@ -10,7 +10,7 @@
 class TestNormalProposalDistribution(chex.TestCase):
     def setUp(self):
         super().setUp()
-        self.key = jax.random.PRNGKey(20220611)
+        self.key = jax.random.key(20220611)
 
     def test_normal_univariate(self):
         """

tests/mcmc/test_random_walk_without_chex.py

@@ -21,7 +21,7 @@ def test_one_step_addition(self):
         Since the density == 1, the proposal is accepted.
         The random step may depend on the previous position
         """
-        rng_key = jax.random.PRNGKey(0)
+        rng_key = jax.random.key(0)
         initial_position = jnp.array([50.0])
 
         def random_step(key, position):
@@ -51,7 +51,7 @@ def test_logdensity_accepts(position):
 class IRMHTest(unittest.TestCase):
     def test_proposal_is_independent_of_position(self):
         """New position does not depend on previous"""
-        rng_key = jax.random.PRNGKey(0)
+        rng_key = jax.random.key(0)
         initial_position = jnp.array([50.0])
         other_position = jnp.array([15000.0])
 
@@ -99,7 +99,7 @@ def test_generate_reject(self):
         and given that the sampling rule rejects,
         the prev_state is proposed again
         """
-        rng_key = jax.random.PRNGKey(0)
+        rng_key = jax.random.key(0)
 
         prev_state = RWState(jnp.array([30.0]), 15.0)
 
@@ -118,7 +118,7 @@ def test_generate_reject(self):
         np.testing.assert_allclose(sampled_proposal.state.position, jnp.array([30.0]))
 
     def test_generate_accept(self):
-        rng_key = jax.random.PRNGKey(0)
+        rng_key = jax.random.key(0)
         prev_state = RWState(jnp.array([30.0]), 15.0)
 
         generate = rmh_proposal(

tests/mcmc/test_sampling.py

@@ -71,7 +71,7 @@ class LinearRegressionTest(chex.TestCase):
 
     def setUp(self):
         super().setUp()
-        self.key = jax.random.PRNGKey(19)
+        self.key = jax.random.key(19)
 
     def regression_logprob(self, log_scale, coefs, preds, x):
         """Linear regression"""
@@ -229,7 +229,7 @@ class SGMCMCTest(chex.TestCase):
 
     def setUp(self):
         super().setUp()
-        self.key = jax.random.PRNGKey(19)
+        self.key = jax.random.key(19)
 
     def logprior_fn(self, position):
         return -0.5 * jnp.dot(position, position) * 0.01
@@ -379,7 +379,7 @@ class LatentGaussianTest(chex.TestCase):
 
     def setUp(self):
         super().setUp()
-        self.key = jax.random.PRNGKey(19)
+        self.key = jax.random.key(19)
         self.C = 2.0 * np.eye(1)
         self.delta = 5.0
         self.sampling_steps = 25_000
@@ -494,7 +494,7 @@ class UnivariateNormalTest(chex.TestCase):
 
     def setUp(self):
         super().setUp()
-        self.key = jax.random.PRNGKey(12)
+        self.key = jax.random.key(12)
 
     def normal_logprob(self, x):
         return stats.norm.logpdf(x, loc=1.0, scale=2.0)
@@ -572,7 +572,7 @@ class MonteCarloStandardErrorTest(chex.TestCase):
 
     def setUp(self):
         super().setUp()
-        self.key = jax.random.PRNGKey(20220203)
+        self.key = jax.random.key(20220203)
 
     def generate_multivariate_target(self, rng=None):
         """Genrate a Multivariate Normal distribution as target."""

tests/mcmc/test_trajectory.py

@@ -22,7 +22,7 @@ class TrajectoryTest(chex.TestCase):
     def test_dynamic_progressive_integration_divergence(
         self, step_size, should_diverge
     ):
-        rng_key = jax.random.PRNGKey(0)
+        rng_key = jax.random.key(0)
 
         logdensity_fn = jax.scipy.stats.norm.logpdf
 
@@ -74,7 +74,7 @@ def test_dynamic_progressive_integration_divergence(
         assert is_diverging.item() is should_diverge
 
     def test_dynamic_progressive_equal_recursive(self):
-        rng_key = jax.random.PRNGKey(23132)
+        rng_key = jax.random.key(23132)
 
         def logdensity_fn(x):
             return -((1.0 - x[0]) ** 2) - 1.5 * (x[1] - x[0] ** 2) ** 2
@@ -202,7 +202,7 @@ def logdensity_fn(x):
     def test_dynamic_progressive_expansion(
         self, step_size, should_diverge, should_turn, expected_doublings
     ):
-        rng_key = jax.random.PRNGKey(0)
+        rng_key = jax.random.key(0)
 
         def logdensity_fn(x):
             return -0.5 * x**2
@@ -260,7 +260,7 @@ def logdensity_fn(x):
         assert is_turning == should_turn
 
     def test_static_integration_variable_num_steps(self):
-        rng_key = jax.random.PRNGKey(0)
+        rng_key = jax.random.key(0)
 
         logdensity_fn = jax.scipy.stats.norm.logpdf
         position = 1.0

tests/optimizers/test_optimizers.py

@@ -23,7 +23,7 @@
 class OptimizerTest(chex.TestCase):
     def setUp(self):
         super().setUp()
-        self.key = jax.random.PRNGKey(1)
+        self.key = jax.random.key(1)
 
     @chex.all_variants(with_pmap=False)
     def test_dual_averaging(self):

tests/optimizers/test_pathfinder.py

@@ -14,7 +14,7 @@
 class PathfinderTest(chex.TestCase):
     def setUp(self):
         super().setUp()
-        self.key = jax.random.PRNGKey(1)
+        self.key = jax.random.key(1)
 
     @chex.all_variants(without_device=False, with_pmap=False)
     @parameterized.parameters(

tests/smc/test_kernel_compatibility.py

@@ -17,7 +17,7 @@ class SMCAndMCMCIntegrationTest(unittest.TestCase):
 
     def setUp(self):
         super().setUp()
-        self.key = jax.random.PRNGKey(42)
+        self.key = jax.random.key(42)
         self.initial_particles = jax.random.multivariate_normal(
             self.key, jnp.zeros(2), jnp.eye(2), (3,)
         )

tests/smc/test_resampling.py

@@ -41,7 +41,7 @@ def test_resampling_methods(self, num_samples, method_name):
         x = jnp.array(np.random.randn(N), dtype="float32")
         w = w / w.sum()
 
-        resampling_keys = jax.random.split(jax.random.PRNGKey(42), batch_size)
+        resampling_keys = jax.random.split(jax.random.key(42), batch_size)
 
         resampling_idx = jax.vmap(
             self.variant(resampling_methods[method_name], static_argnums=(2,)),

tests/smc/test_smc.py

@@ -24,7 +24,7 @@ def _weighted_avg_and_std(values, weights):
 class SMCTest(chex.TestCase):
     def setUp(self):
         super().setUp()
-        self.key = jax.random.PRNGKey(42)
+        self.key = jax.random.key(42)
 
     @chex.variants(with_jit=True)
     def test_smc(self):

tests/smc/test_tempered_smc.py

@@ -37,7 +37,7 @@ class TemperedSMCTest(chex.TestCase):
 
     def setUp(self):
         super().setUp()
-        self.key = jax.random.PRNGKey(42)
+        self.key = jax.random.key(42)
 
     def logdensity_fn(self, log_scale, coefs, preds, x):
         """Linear regression"""
@@ -174,7 +174,7 @@ def test_normalizing_constant(self):
         num_particles = 200
         num_dim = 2
 
-        rng_key = jax.random.PRNGKey(2356)
+        rng_key = jax.random.key(2356)
         rng_key, cov_key = jax.random.split(rng_key, 2)
         chol_cov = jax.random.uniform(cov_key, shape=(num_dim, num_dim))
         iu = np.triu_indices(num_dim, 1)

tests/test_benchmarks.py

@@ -38,7 +38,7 @@ def one_step(state, rng_key):
 
 
 def run_regression(algorithm, **parameters):
-    key = jax.random.PRNGKey(0)
+    key = jax.random.key(0)
     rng_key, init_key0, init_key1 = jax.random.split(key, 3)
     x_data = jax.random.normal(init_key0, shape=(100_000, 1))
     y_data = 3 * x_data + jax.random.normal(init_key1, shape=x_data.shape)