address expanded distribution

Author: fehiepsi

numpyro/handlers.py

@@ -83,7 +83,7 @@
 import jax.numpy as jnp
 
 import numpyro
-from numpyro.distributions.distribution import COERCIONS
+from numpyro.distributions.distribution import COERCIONS, ExpandedDistribution
 from numpyro.primitives import _PYRO_STACK, Messenger, apply_stack, plate
 from numpyro.util import not_jax_tracer
 
@@ -268,6 +268,8 @@ def process_message(self, msg):
         if msg["type"] == "sample":
             if msg["value"] is None:
                 msg["value"] = msg["name"]
+                if isinstance(msg["fn"], ExpandedDistribution):
+                    msg["fn"] = msg["fn"].base_dist
 
             if isinstance(msg["fn"], Funsor) or isinstance(msg["value"], (str, Funsor)):
                 msg["stop"] = True

test/test_handlers.py

@@ -636,7 +636,6 @@ def model():
         x = numpyro.sample("x", dist.Normal(0, 1))
         with handlers.collapse():
             with handlers.plate("data", len(data)):
-                # TODO: address expanded distribution
                 y = numpyro.sample("y", dist.Normal(x, 1.))
                 numpyro.sample("z", dist.Normal(y, 1.), obs=data)