make expand works under funsor

Author: fehiepsi

numpyro/handlers.py

@@ -77,13 +77,14 @@
 """
 
 from collections import OrderedDict
+from functools import reduce
 import warnings
 
-from jax import lax, random
+from jax import lax, random, tree_map
 import jax.numpy as jnp
 
 import numpyro
-from numpyro.distributions.distribution import COERCIONS, ExpandedDistribution
+from numpyro.distributions.distribution import COERCIONS, ExpandedDistribution, Independent
 from numpyro.primitives import _PYRO_STACK, Messenger, apply_stack, plate
 from numpyro.util import not_jax_tracer
 
@@ -245,6 +246,20 @@ def process_message(self, msg):
             msg['stop'] = True
 
 
+def _eager_expand_fn(fn):
+    if isinstance(fn, Independent):
+        reinterpreted_batch_ndims = fn.reinterpreted_batch_ndims
+        fn = fn.base_dist
+    else:
+        reinterpreted_batch_ndims = 0  # no-op for to_event method
+    if isinstance(fn, ExpandedDistribution):
+        batch_shape = fn.batch_shape
+        base_batch_shape = fn.base_dist.batch_shape
+        appended_shape = batch_shape[:len(batch_shape) - len(base_batch_shape)]
+        fn = tree_map(lambda x: jnp.broadcast_to(x, appended_shape + jnp.shape(x)), fn.base_dist)
+    return fn.to_event(reinterpreted_batch_ndims)
+
+
 class collapse(trace):
     """
     EXPERIMENTAL Collapses all sites in the context by lazily sampling and
@@ -263,33 +278,48 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
     def process_message(self, msg):
-        from funsor.terms import Funsor
+        if msg["type"] != "sample":
+            return
 
-        if msg["type"] == "sample":
-            if msg["value"] is None:
-                msg["value"] = msg["name"]
-                if isinstance(msg["fn"], ExpandedDistribution):
-                    msg["fn"] = msg["fn"].base_dist
+        import funsor
+
+        # Eagerly convert fn and value to Funsor.
+        dim_to_name = {f.dim: f.name for f in msg["cond_indep_stack"]}
+        dim_to_name.update(self.preserved_plates)
+        if isinstance(msg["fn"], (Independent, ExpandedDistribution)):
+            msg["fn"] = _eager_expand_fn(msg["fn"])
+        msg["fn"] = funsor.to_funsor(msg["fn"], funsor.Real, dim_to_name)
+        domain = msg["fn"].inputs["value"]
+        if msg["value"] is None:
+            msg["value"] = funsor.Variable(msg["name"], domain)
+        else:
+            msg["value"] = funsor.to_funsor(msg["value"], domain, dim_to_name)
 
-            if isinstance(msg["fn"], Funsor) or isinstance(msg["value"], (str, Funsor)):
-                msg["stop"] = True
+        msg["stop"] = True
 
     def __enter__(self):
-        self.preserved_plates = frozenset(h.name for h in _PYRO_STACK
-                                          if isinstance(h, plate))
+        self.preserved_plates = {h.dim: h.name for h in _PYRO_STACK
+                                 if isinstance(h, plate)}
         COERCIONS.append(self._coerce)
         return super().__enter__()
 
     def __exit__(self, exc_type, exc_value, traceback):
-        import funsor
-
         _coerce = COERCIONS.pop()
         assert _coerce is self._coerce
         super().__exit__(exc_type, exc_value, traceback)
 
         if exc_type is not None:
+            self.trace.clear()
+            self.preserved_plates.clear()
             return
 
+        if any(site["type"] == "sample" for site in self.trace.values()):
+            name, log_prob, _, _ = self._get_log_prob()
+            numpyro.factor(name, log_prob.data)
+
+    def _get_log_prob(self):
+        import funsor
+
         # Convert delayed statements to pyro.factor()
         reduced_vars = []
         log_prob_terms = []
@@ -299,24 +329,28 @@ def __exit__(self, exc_type, exc_value, traceback):
                 continue
             if not site["is_observed"]:
                 reduced_vars.append(name)
-            dim_to_name = {f.dim: f.name for f in site["cond_indep_stack"]}
-            fn = funsor.to_funsor(site["fn"], funsor.Real, dim_to_name)
-            value = site["value"]
-            if not isinstance(value, str):
-                value = funsor.to_funsor(site["value"], fn.inputs["value"], dim_to_name)
-            log_prob_terms.append(fn(value=value))
+            log_prob_terms.append(site["fn"](value=site["value"]))
             plates |= frozenset(f.name for f in site["cond_indep_stack"])
-        assert log_prob_terms, "nothing to collapse"
-        reduced_plates = plates - self.preserved_plates
-        log_prob = funsor.sum_product.sum_product(
-            funsor.ops.logaddexp,
-            funsor.ops.add,
-            log_prob_terms,
-            eliminate=frozenset(reduced_vars) | reduced_plates,
-            plates=plates,
-        )
         name = reduced_vars[0]
-        numpyro.factor(name, log_prob.data)
+        reduced_vars = frozenset(reduced_vars)
+        assert log_prob_terms, "nothing to collapse"
+        reduced_plates = plates - frozenset(self.preserved_plates.values())
+        self.trace.clear()
+        self.preserved_plates.clear()
+        if reduced_plates:
+            log_prob = funsor.sum_product.sum_product(
+                funsor.ops.logaddexp,
+                funsor.ops.add,
+                log_prob_terms,
+                eliminate=frozenset(reduced_vars) | reduced_plates,
+                plates=plates,
+            )
+            log_joint = NotImplemented
+        else:
+            log_joint = reduce(funsor.ops.add, log_prob_terms)
+            log_prob = log_joint.reduce(funsor.ops.logaddexp, reduced_vars)
+
+        return name, log_prob, log_joint, reduced_vars
 
 
 class condition(Messenger):

test/test_handlers.py

@@ -656,15 +656,14 @@ def test_collapse_normal_mvn_mvn():
     def model():
         x = numpyro.sample("x", dist.Exponential(1))
         with handlers.collapse():
-            with numpyro.plate("d", d):
-                # TODO: verify that to_event works here
-                beta0 = numpyro.sample("beta0", dist.Normal(0, 1).expand([S]).to_event(1))
-                # TODO: address beta0 is a str, which cannot do infer_param_domain
-                beta = numpyro.sample("beta", dist.MultivariateNormal(beta0, jnp.eye(S)))
-            # FIXME: beta is a string here, how to apply numeric operators
+            with numpyro.plate("d", d, dim=-1):
+                beta0 = numpyro.sample("beta0", dist.Normal(0., 1.).expand([d, S]).to_event(1))
+                beta = numpyro.sample(
+                    "beta", dist.MultivariateNormal(beta0, scale_tril=jnp.eye(S)))
+
             mean = jnp.ones((T, d)) @ beta
-            with numpyro.plate("data", T, dim=-2):
-                numpyro.sample("obs", dist.MultivariateNormal(mean, jnp.eye(S)), obs=data)
+            with numpyro.plate("data", T, dim=-1):
+                numpyro.sample("obs", dist.MultivariateNormal(mean, scale_tril=jnp.eye(S)), obs=data)
 
     def guide():
         rate = numpyro.param("rate", 1., constraint=constraints.positive)