Repository: pyro-ppl/numpyro
Pull request #1823 (2024-06-30 20:23:33)

Fixes `random_flax_module` with `flax.linen.BatchNorm`
Fixes https://github.com/pyro-ppl/numpyro/issues/1446

juanitorduz: The first two commits are from an old branch 🤦 . We can squash and merge instead.
juanitorduz: https://preview.redd.it/sometimes-progress-looks-like-failure-v0-x4yghusfc1o81.jpg?width=640&crop=smart&auto=webp&s=df5a49d386fca48eaeffaea8945f2060782fcfbc
juanitorduz: ![image](https://github.com/pyro-ppl/numpyro/assets/22996444/d369d911-d1ba-403c-b7e7-65e610ed5de8)

juanitorduz: @fehiepsi should I leave `_substitute_default_key` in `utils.py` or is `handlers.py` a better place? 
fehiepsi: Leaving it in `utils` sounds reasonable to me. It is just a workaround for the edge case.

Thanks for fixing the issue!

Patch:
diff --git a/numpyro/handlers.py b/numpyro/handlers.py
index 6e13aeb7..11fdde45 100644
--- a/numpyro/handlers.py
+++ b/numpyro/handlers.py
@@ -817,7 +817,7 @@ class substitute(Messenger):
             return
 
         if self.data is not None:
-            value = self.data.get(msg["name"])
+            value = self.data.get(msg.get("name"))
         else:
             value = self.substitute_fn(msg)
 
diff --git a/numpyro/infer/util.py b/numpyro/infer/util.py
index 3775893a..56c4d2c4 100644
--- a/numpyro/infer/util.py
+++ b/numpyro/infer/util.py
@@ -21,6 +21,7 @@ from numpyro.distributions.transforms import biject_to
 from numpyro.distributions.util import is_identically_one, sum_rightmost
 from numpyro.handlers import condition, replay, seed, substitute, trace
 from numpyro.infer.initialization import init_to_uniform, init_to_value
+from numpyro.primitives import Messenger
 from numpyro.util import (
     _validate_model,
     find_stack_level,
@@ -46,6 +47,12 @@ ModelInfo = namedtuple(
 ParamInfo = namedtuple("ParamInfo", ["z", "potential_energy", "z_grad"])
 
 
+class _substitute_default_key(Messenger):
+    def process_message(self, msg):
+        if msg["type"] == "prng_key" and msg["value"] is None:
+            msg["value"] = random.PRNGKey(0)
+
+
 def log_density(model, model_args, model_kwargs, params):
     """
     (EXPERIMENTAL INTERFACE) Computes log of joint density for the model given
@@ -660,9 +667,12 @@ def initialize_model(
         data={
             k: site["value"]
             for k, site in model_trace.items()
-            if site["type"] in ["param"]
+            if site["type"] in ["param", "mutable"]
         },
     )
+
+    model = _substitute_default_key(model)
+
     constrained_values = {
         k: v["value"]
         for k, v in model_trace.items()


Test patch:
diff --git a/test/contrib/test_module.py b/test/contrib/test_module.py
index d3f27f17..a1342507 100644
--- a/test/contrib/test_module.py
+++ b/test/contrib/test_module.py
@@ -21,7 +21,8 @@ from numpyro.contrib.module import (
     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 test_haiku_state_dropout_smoke(dropout, batchnorm):
     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 test_flax_state_dropout_smoke(dropout, batchnorm):
         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)