Fixes random_flax_module with flax.linen.BatchNorm

numpyro/handlers.py

@@ -805,7 +805,8 @@ def __init__(self, fn=None, data=None, substitute_fn=None):
 
     def process_message(self, msg):
         if (
-            msg["type"] not in ("sample", "param", "mutable", "plate", "deterministic")
+            msg["type"]
+            not in ("sample", "param", "mutable", "plate", "deterministic", "prng_key")
         ) or msg.get("_control_flow_done", False):
             if msg["type"] == "control_flow":
                 if self.data is not None:

numpyro/infer/util.py

@@ -653,16 +653,17 @@ def initialize_model(
         has_enumerate_support,
         model_trace,
     ) = _get_model_transforms(substituted_model, model_args, model_kwargs)
+
     # substitute param sites from model_trace to model so
     # we don't need to generate again parameters of `numpyro.module`
-    model = substitute(
-        model,
-        data={
-            k: site["value"]
-            for k, site in model_trace.items()
-            if site["type"] in ["param"]
-        },
-    )
+    def substitute_fn(site):
+        if site["type"] in ["param", "mutable"]:
+            return site["value"]
+        elif site["type"] == "prng_key":
+            return random.PRNGKey(0)
+
+    model = substitute(model, substitute_fn=substitute_fn)
+
     constrained_values = {
         k: v["value"]
         for k, v in model_trace.items()
@@ -701,14 +702,7 @@ def initialize_model(
     prototype_params = transform_fn(inv_transforms, constrained_values, invert=True)
     (init_params, pe, grad), is_valid = find_valid_initial_params(
         rng_key,
-        substitute(
-            model,
-            data={
-                k: site["value"]
-                for k, site in model_trace.items()
-                if site["type"] in ["plate"]
-            },
-        ),
+        substitute(model, substitute_fn=substitute_fn),
         init_strategy=init_strategy,
         enum=has_enumerate_support,
         model_args=model_args,

test/contrib/test_module.py

@@ -21,7 +21,8 @@
     random_haiku_module,
 )
 import numpyro.distributions as dist
-from numpyro.infer import MCMC, NUTS
+from numpyro.infer import MCMC, NUTS, SVI, Trace_ELBO
+from numpyro.infer.autoguide import AutoDelta
 
 pytestmark = pytest.mark.filterwarnings(
     "ignore:jax.tree_.+ is deprecated:FutureWarning"
@@ -256,6 +257,11 @@ def model():
     else:
         assert set(tr.keys()) == {"nn$params", "x", "y"}
 
+    # test svi
+    guide = AutoDelta(model)
+    svi = SVI(model, guide, numpyro.optim.Adam(0.01), Trace_ELBO())
+    svi.run(random.PRNGKey(100), 10)
+
 
 @pytest.mark.parametrize("dropout", [True, False])
 @pytest.mark.parametrize("batchnorm", [True, False])
@@ -300,3 +306,8 @@ def model():
         assert tr["nn$state"]["type"] == "mutable"
     else:
         assert set(tr.keys()) == {"nn$params", "x", "y"}
+
+    # test svi
+    guide = AutoDelta(model)
+    svi = SVI(model, guide, numpyro.optim.Adam(0.01), Trace_ELBO())
+    svi.run(random.PRNGKey(100), 10)