Repository: pyro-ppl/numpyro
Pull request #1992 (2025-03-01 13:26:42)

Improve rng default in the `contrib/model.py` module
Closes https://github.com/pyro-ppl/numpyro/issues/1991

fehiepsi: I think it is the issue of contrib.module. We can use
```
rng_key = numpyro.prng_key() or random.key(0)
```
to optionally create dummy rng_key. I think it does not affect the functionality of the module.
juanitorduz: > 5e40979

Do you mean like https://github.com/pyro-ppl/numpyro/pull/1992/commits/5e409799f8f42aaed4520d8b2a14d0d4c697d2d9 ? or something has to be changed in `contrib/module.py`?
juanitorduz: > I think it is the issue of contrib.module. We can use
> 
> ```
> rng_key = numpyro.prng_key() or random.key(0)
> ```
> 
> to optionally create dummy rng_key. I think it does not affect the functionality of the module.

Do you mean like https://github.com/pyro-ppl/numpyro/pull/1992/commits/5e409799f8f42aaed4520d8b2a14d0d4c697d2d9 ? or something has to be changed in `contrib/module.py`?
fehiepsi: I think changing in contrib is better because we know that key(0) does not change the logic there. If we change it here, users might forget to provide some posterior samples.
fehiepsi: Could you merge with the master branch?

Patch:
diff --git a/numpyro/contrib/module.py b/numpyro/contrib/module.py
index f87d96d5..a846373f 100644
--- a/numpyro/contrib/module.py
+++ b/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 patch:
diff --git a/test/infer/test_infer_util.py b/test/infer/test_infer_util.py
index ab013370..c9ee9e60 100644
--- a/test/infer/test_infer_util.py
+++ b/test/infer/test_infer_util.py
@@ -519,3 +519,44 @@ def test_get_mask_optimization():
     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)