Improve rng default in the contrib/model.py module (#1992)

* fix

* rm comment

* alternative rng

* suggestion

Author: Juan Orduz

numpyro/contrib/module.py

@@ -85,6 +85,8 @@ def flax_module(
         # feed in dummy data to init params
         args = (jnp.ones(input_shape),) if input_shape is not None else args
         rng_key = numpyro.prng_key()
+        if rng_key is None:
+            rng_key = random.key(0)
         # split rng_key into a dict of rng_kind: rng_key
         rngs = {}
         if apply_rng:
@@ -187,6 +189,8 @@ def haiku_module(name, nn_module, *args, input_shape=None, apply_rng=False, **kw
         args = (jnp.ones(input_shape),) if input_shape is not None else args
         # feed in dummy data to init params
         rng_key = numpyro.prng_key()
+        if rng_key is None:
+            rng_key = random.key(0)
         if with_state:
             nn_params, nn_state = nn_module.init(rng_key, *args, **kwargs)
             nn_state = dict(nn_state)

test/infer/test_infer_util.py

@@ -519,3 +519,44 @@ def guide():
     assert "guide-always" in called
     assert "model-sometimes" not in called
     assert "guide-sometimes" not in called
+
+
+def test_log_likelihood_flax_nn():
+    import numpy as np
+
+    import flax.linen as nn
+    from jax import random
+
+    from numpyro.contrib.module import random_flax_module
+
+    # Simulate
+    rng = np.random.default_rng(99)
+    N = 1000
+
+    X = rng.normal(0, 1, size=(N, 1))
+    mu = 1 + X @ np.array([0.5])
+    y = rng.normal(mu, 0.5)
+
+    # Simple linear layer
+    class Linear(nn.Module):
+        @nn.compact
+        def __call__(self, x):
+            return nn.Dense(1, use_bias=True, name="Dense")(x)
+
+    def model(X, y=None):
+        sigma = numpyro.sample("sigma", dist.HalfNormal(0.1))
+        priors = {"Dense.bias": dist.Normal(0, 2.5), "Dense.kernel": dist.Normal(0, 1)}
+        mlp = random_flax_module(
+            "mlp", Linear(), prior=priors, input_shape=(X.shape[1],)
+        )
+        with numpyro.plate("data", X.shape[0]):
+            mu = numpyro.deterministic("mu", mlp(X).squeeze(-1))
+            y = numpyro.sample("y", dist.Normal(mu, sigma), obs=y)
+
+    # Fit model
+    kernel = NUTS(model, target_accept_prob=0.95)
+    mcmc = MCMC(kernel, num_warmup=100, num_samples=100, num_chains=1)
+    mcmc.run(random.PRNGKey(0), X=X, y=y)
+
+    # run log likelihood
+    numpyro.infer.util.log_likelihood(model, mcmc.get_samples(), X=X, y=y)