Support custom prng key (#1642)

* support custom prng key

* run black

* test custom prng in CI

* fix some deprecation warnings

Author: fehiepsi

.github/workflows/ci.yml

@@ -112,6 +112,9 @@ jobs:
         XLA_FLAGS="--xla_force_host_platform_device_count=2" pytest -vs test/infer/test_mcmc.py -k "chain or pmap or vmap"
         XLA_FLAGS="--xla_force_host_platform_device_count=2" pytest -vs test/contrib/test_tfp.py -k "chain"
         XLA_FLAGS="--xla_force_host_platform_device_count=2" pytest -vs test/infer/test_hmc_gibbs.py -k "chain"
+    - name: Test custom prng
+      run: |
+        JAX_ENABLE_CUSTOM_PRNG=1 pytest -vs test/infer/test_mcmc.py
 
 
   examples:

numpyro/contrib/einstein/steinvi.py

@@ -10,7 +10,6 @@
 import operator
 
 from jax import grad, jacfwd, numpy as jnp, random, vmap
-from jax.random import KeyArray
 from jax.tree_util import tree_map
 
 from numpyro import handlers
@@ -370,10 +369,10 @@ def _update_force(attr_force, rep_force, jac):
         )
         return jnp.linalg.norm(particle_grads), res_grads
 
-    def init(self, rng_key: KeyArray, *args, **kwargs):
+    def init(self, rng_key, *args, **kwargs):
         """Register random variable transformations, constraints and determine initialize positions of the particles.
 
-        :param KeyArray rng_key: Random number generator seed.
+        :param rng_key: Random number generator seed.
         :param args: Arguments to the model / guide.
         :param kwargs: Keyword arguments to the model / guide.
         :return: initial :data:`SteinVIState`

numpyro/contrib/tfp/mcmc.py

@@ -14,7 +14,7 @@
 from numpyro.infer import init_to_uniform
 from numpyro.infer.mcmc import MCMCKernel
 from numpyro.infer.util import initialize_model
-from numpyro.util import identity
+from numpyro.util import identity, is_prng_key
 
 TFPKernelState = namedtuple("TFPKernelState", ["z", "kernel_results", "rng_key"])
 
@@ -174,7 +174,7 @@ def init(
         self, rng_key, num_warmup, init_params=None, model_args=(), model_kwargs={}
     ):
         # non-vectorized
-        if rng_key.ndim == 1:
+        if is_prng_key(rng_key):
             rng_key, rng_key_init_model = random.split(rng_key)
         # vectorized
         else:
@@ -190,7 +190,7 @@ def init(
                 " `target_log_prob_fn`."
             )
 
-        if rng_key.ndim == 1:
+        if is_prng_key(rng_key):
             init_state = self._init_fn(init_params, rng_key)
         else:
             # XXX it is safe to run hmc_init_fn under vmap despite that hmc_init_fn changes some

numpyro/distributions/conjugate.py

@@ -14,7 +14,8 @@
     ZeroInflatedDistribution,
 )
 from numpyro.distributions.distribution import Distribution
-from numpyro.distributions.util import is_prng_key, promote_shapes, validate_sample
+from numpyro.distributions.util import promote_shapes, validate_sample
+from numpyro.util import is_prng_key
 
 
 def _log_beta_1(alpha, value):

numpyro/distributions/continuous.py

@@ -64,14 +64,14 @@
     betaincinv,
     cholesky_of_inverse,
     gammaincinv,
-    is_prng_key,
     lazy_property,
     matrix_to_tril_vec,
     promote_shapes,
     signed_stick_breaking_tril,
     validate_sample,
     vec_to_tril_matrix,
 )
+from numpyro.util import is_prng_key
 
 
 class AsymmetricLaplace(Distribution):

numpyro/distributions/copula.py

@@ -6,12 +6,8 @@
 import numpyro.distributions.constraints as constraints
 from numpyro.distributions.continuous import Beta, MultivariateNormal, Normal
 from numpyro.distributions.distribution import Distribution
-from numpyro.distributions.util import (
-    clamp_probs,
-    is_prng_key,
-    lazy_property,
-    validate_sample,
-)
+from numpyro.distributions.util import clamp_probs, lazy_property, validate_sample
+from numpyro.util import is_prng_key
 
 
 class GaussianCopula(Distribution):

numpyro/distributions/directional.py

@@ -16,14 +16,13 @@
 from numpyro.distributions import constraints
 from numpyro.distributions.distribution import Distribution
 from numpyro.distributions.util import (
-    is_prng_key,
     lazy_property,
     promote_shapes,
     safe_normalize,
     validate_sample,
     von_mises_centered,
 )
-from numpyro.util import while_loop
+from numpyro.util import is_prng_key, while_loop
 
 
 def _numel(shape):

numpyro/distributions/discrete.py

@@ -41,13 +41,12 @@
     binomial,
     categorical,
     clamp_probs,
-    is_prng_key,
     lazy_property,
     multinomial,
     promote_shapes,
     validate_sample,
 )
-from numpyro.util import not_jax_tracer
+from numpyro.util import is_prng_key, not_jax_tracer
 
 
 def _to_probs_bernoulli(logits):

numpyro/distributions/mixtures.py

@@ -7,7 +7,8 @@
 
 from numpyro.distributions import Distribution, constraints
 from numpyro.distributions.discrete import CategoricalLogits, CategoricalProbs
-from numpyro.distributions.util import is_prng_key, validate_sample
+from numpyro.distributions.util import validate_sample
+from numpyro.util import is_prng_key
 
 
 def Mixture(mixing_distribution, component_distributions, *, validate_args=None):

numpyro/distributions/truncated.py

@@ -19,11 +19,11 @@
 from numpyro.distributions.distribution import Distribution
 from numpyro.distributions.util import (
     clamp_probs,
-    is_prng_key,
     lazy_property,
     promote_shapes,
     validate_sample,
 )
+from numpyro.util import is_prng_key
 
 
 class LeftTruncatedDistribution(Distribution):