Add cond primitive (#1028)

tcbegley · web-flow · commit 1ca72ffc3ddd · 2021-05-04T08:53:31.000-05:00
* Add `cond` control flow primitive

* Add test for cond

* Make cond usage example more readable

* Docstring fixes

* Test NUTS with cond

* Add return statements to example

* Fix markup in docstring

* Fix failing doctest

* Add enumeration warning

* Make MCMC + cond test more robust

* Test full mean and std of mixture in cond test

* Adjust mode separation in cond test

* Typo in test &gt;_&lt;
diff --git a/docs/source/primitives.rst b/docs/source/primitives.rst
@@ -62,3 +62,7 @@ random_haiku_module
 scan
 ----
 .. autofunction:: numpyro.contrib.control_flow.scan
+
+cond
+----
+.. autofunction:: numpyro.contrib.control_flow.cond
diff --git a/numpyro/contrib/control_flow/__init__.py b/numpyro/contrib/control_flow/__init__.py
@@ -1,8 +1,7 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
 
+from numpyro.contrib.control_flow.cond import cond
 from numpyro.contrib.control_flow.scan import scan
 
-__all__ = [
-    "scan",
-]
+__all__ = ["cond", "scan"]
diff --git a/numpyro/contrib/control_flow/cond.py b/numpyro/contrib/control_flow/cond.py
@@ -0,0 +1,162 @@
+# Copyright Contributors to the Pyro project.
+# SPDX-License-Identifier: Apache-2.0
+
+from functools import partial
+
+from jax import device_put, lax
+
+from numpyro import handlers
+from numpyro.contrib.control_flow.util import PytreeTrace
+from numpyro.primitives import _PYRO_STACK, apply_stack
+
+
+def _subs_wrapper(subs_map, site):
+    if isinstance(subs_map, dict) and site["name"] in subs_map:
+        return subs_map[site["name"]]
+    elif callable(subs_map):
+        rng_key = site["kwargs"].get("rng_key")
+        subs_map = (
+            handlers.seed(subs_map, rng_seed=rng_key)
+            if rng_key is not None
+            else subs_map
+        )
+        return subs_map(site)
+    return None
+
+
+def wrap_fn(fn, substitute_stack):
+    def wrapper(wrapped_operand):
+        rng_key, operand = wrapped_operand
+
+        with handlers.block():
+            seeded_fn = handlers.seed(fn, rng_key) if rng_key is not None else fn
+            for subs_type, subs_map in substitute_stack:
+                subs_fn = partial(_subs_wrapper, subs_map)
+                if subs_type == "condition":
+                    seeded_fn = handlers.condition(seeded_fn, condition_fn=subs_fn)
+                elif subs_type == "substitute":
+                    seeded_fn = handlers.substitute(seeded_fn, substitute_fn=subs_fn)
+
+            with handlers.trace() as trace:
+                value = seeded_fn(operand)
+
+        return value, PytreeTrace(trace)
+
+    return wrapper
+
+
+def cond_wrapper(
+    pred,
+    true_fun,
+    false_fun,
+    operand,
+    rng_key=None,
+    substitute_stack=None,
+    enum=False,
+    first_available_dim=None,
+):
+    if enum:
+        # TODO: support enumeration. note that pred passed to lax.cond must be scalar
+        # which means that even simple conditions like `x == 0` can get complicated if
+        # x is an enumerated discrete random variable
+        raise RuntimeError("The cond primitive does not currently support enumeration")
+
+    if substitute_stack is None:
+        substitute_stack = []
+
+    wrapped_true_fun = wrap_fn(true_fun, substitute_stack)
+    wrapped_false_fun = wrap_fn(false_fun, substitute_stack)
+    wrapped_operand = device_put((rng_key, operand))
+    return lax.cond(pred, wrapped_true_fun, wrapped_false_fun, wrapped_operand)
+
+
+def cond(pred, true_fun, false_fun, operand):
+    """
+    This primitive conditionally applies ``true_fun`` or ``false_fun``. See
+    :func:`jax.lax.cond` for more information.
+
+    **Usage**:
+
+    .. doctest::
+
+       >>> import numpyro
+       >>> import numpyro.distributions as dist
+       >>> from jax import random
+       >>> from numpyro.contrib.control_flow import cond
+       >>> from numpyro.infer import SVI, Trace_ELBO
+       >>>
+       >>> def model():
+       ...     def true_fun(_):
+       ...         return numpyro.sample("x", dist.Normal(20.0))
+       ...
+       ...     def false_fun(_):
+       ...         return numpyro.sample("x", dist.Normal(0.0))
+       ...
+       ...     cluster = numpyro.sample("cluster", dist.Normal())
+       ...     return cond(cluster > 0, true_fun, false_fun, None)
+       >>>
+       >>> def guide():
+       ...     m1 = numpyro.param("m1", 10.0)
+       ...     s1 = numpyro.param("s1", 0.1, constraint=dist.constraints.positive)
+       ...     m2 = numpyro.param("m2", 10.0)
+       ...     s2 = numpyro.param("s2", 0.1, constraint=dist.constraints.positive)
+       ...
+       ...     def true_fun(_):
+       ...         return numpyro.sample("x", dist.Normal(m1, s1))
+       ...
+       ...     def false_fun(_):
+       ...         return numpyro.sample("x", dist.Normal(m2, s2))
+       ...
+       ...     cluster = numpyro.sample("cluster", dist.Normal())
+       ...     return cond(cluster > 0, true_fun, false_fun, None)
+       >>>
+       >>> svi = SVI(model, guide, numpyro.optim.Adam(1e-2), Trace_ELBO(num_particles=100))
+       >>> params, losses = svi.run(random.PRNGKey(0), num_steps=2500)
+
+    .. warning:: This is an experimental utility function that allows users to use
+        JAX control flow with NumPyro's effect handlers. Currently, `sample` and
+        `deterministic` sites within `true_fun` and `false_fun` are supported. If you
+        notice that any effect handlers or distributions are unsupported, please file
+        an issue.
+
+    .. warning:: The ``cond`` primitive does not currently support enumeration and can
+        not be used inside a ``numpyro.plate`` context.
+
+    .. note:: All ``sample`` sites must belong to the same distribution class. For
+        example the following is not supported
+
+        .. code-block:: python
+
+            cond(
+                True,
+                lambda _: numpyro.sample("x", dist.Normal()),
+                lambda _: numpyro.sample("x", dist.Laplace()),
+                None,
+            )
+
+    :param bool pred: Boolean scalar type indicating which branch function to apply
+    :param callable true_fun: A function to be applied if ``pred`` is true.
+    :param callable false_fun: A function to be applied if ``pred`` is false.
+    :param operand: Operand input to either branch depending on ``pred``. This can
+        be any JAX PyTree (e.g. list / dict of arrays).
+    :return: Output of the applied branch function.
+    """
+    if not _PYRO_STACK:
+        value, _ = cond_wrapper(pred, true_fun, false_fun, operand)
+        return value
+
+    initial_msg = {
+        "type": "control_flow",
+        "fn": cond_wrapper,
+        "args": (pred, true_fun, false_fun, operand),
+        "kwargs": {"rng_key": None, "substitute_stack": []},
+        "value": None,
+    }
+
+    msg = apply_stack(initial_msg)
+    value, pytree_trace = msg["value"]
+
+    for msg in pytree_trace.trace.values():
+        apply_stack(msg)
+
+    return value
diff --git a/numpyro/contrib/control_flow/scan.py b/numpyro/contrib/control_flow/scan.py
@@ -14,45 +14,13 @@
     tree_unflatten,
 )
 import jax.numpy as jnp
-from jax.tree_util import register_pytree_node_class
 
 from numpyro import handlers
+from numpyro.contrib.control_flow.util import PytreeTrace
 from numpyro.primitives import _PYRO_STACK, Messenger, apply_stack
 from numpyro.util import not_jax_tracer
 
 
-@register_pytree_node_class
-class PytreeTrace:
-    def __init__(self, trace):
-        self.trace = trace
-
-    def tree_flatten(self):
-        trace, aux_trace = {}, {}
-        for name, site in self.trace.items():
-            if site["type"] in ["sample", "deterministic"]:
-                trace[name], aux_trace[name] = {}, {"_control_flow_done": True}
-                for key in site:
-                    if key in ["fn", "args", "value", "intermediates"]:
-                        trace[name][key] = site[key]
-                    # scanned sites have stop field because we trace them inside a block handler
-                    elif key != "stop":
-                        aux_trace[name][key] = site[key]
-        # keep the site order information because in JAX, flatten and unflatten do not preserve
-        # the order of keys in a dict
-        site_names = list(trace.keys())
-        return (trace,), (aux_trace, site_names)
-
-    @classmethod
-    def tree_unflatten(cls, aux_data, children):
-        aux_trace, site_names = aux_data
-        (trace,) = children
-        trace_with_aux = {}
-        for name in site_names:
-            trace[name].update(aux_trace[name])
-            trace_with_aux[name] = trace[name]
-        return cls(trace_with_aux)
-
-
 def _subs_wrapper(subs_map, i, length, site):
     value = None
     if isinstance(subs_map, dict) and site["name"] in subs_map:
diff --git a/numpyro/contrib/control_flow/util.py b/numpyro/contrib/control_flow/util.py
@@ -0,0 +1,45 @@
+# Copyright Contributors to the Pyro project.
+# SPDX-License-Identifier: Apache-2.0
+
+from jax.tree_util import register_pytree_node_class
+
+
+@register_pytree_node_class
+class PytreeTrace:
+    def __init__(self, trace):
+        self.trace = trace
+
+    def tree_flatten(self):
+        trace, aux_trace = {}, {}
+        for name, site in self.trace.items():
+            if site["type"] in ["sample", "deterministic"]:
+                trace[name], aux_trace[name] = {}, {"_control_flow_done": True}
+                for key in site:
+                    if key in ["fn", "args", "value", "intermediates"]:
+                        trace[name][key] = site[key]
+                    # scanned sites have stop field because we trace them inside a block handler
+                    elif key != "stop":
+                        if key == "kwargs":
+                            kwargs = site["kwargs"].copy()
+                            if "rng_key" in kwargs:
+                                # rng_key is not traced else it is collected by the
+                                # scan primitive which doesn't make sense
+                                # set to None to avoid leaks during tracing by JAX
+                                kwargs["rng_key"] = None
+                            aux_trace[name][key] = kwargs
+                        else:
+                            aux_trace[name][key] = site[key]
+        # keep the site order information because in JAX, flatten and unflatten do not preserve
+        # the order of keys in a dict
+        site_names = list(trace.keys())
+        return (trace,), (aux_trace, site_names)
+
+    @classmethod
+    def tree_unflatten(cls, aux_data, children):
+        aux_trace, site_names = aux_data
+        (trace,) = children
+        trace_with_aux = {}
+        for name in site_names:
+            trace[name].update(aux_trace[name])
+            trace_with_aux[name] = trace[name]
+        return cls(trace_with_aux)
diff --git a/test/contrib/test_control_flow.py b/test/contrib/test_control_flow.py
@@ -8,10 +8,10 @@
 import jax.numpy as jnp
 
 import numpyro
-from numpyro.contrib.control_flow.scan import scan
+from numpyro.contrib.control_flow import cond, scan
 import numpyro.distributions as dist
 from numpyro.handlers import seed, substitute, trace
-from numpyro.infer import MCMC, NUTS, Predictive
+from numpyro.infer import MCMC, NUTS, SVI, Predictive, Trace_ELBO
 from numpyro.infer.util import potential_energy
 
 
@@ -132,3 +132,72 @@ def iteration(c_prev, c_in):
         result,
         [[1.7, 0.3]],
     )
+
+
+def test_cond():
+    def model():
+        def true_fun(_):
+            x = numpyro.sample("x", dist.Normal(4.0))
+            numpyro.deterministic("z", x - 4.0)
+
+        def false_fun(_):
+            x = numpyro.sample("x", dist.Normal(0.0))
+            numpyro.deterministic("z", x)
+
+        cluster = numpyro.sample("cluster", dist.Normal())
+        cond(cluster > 0, true_fun, false_fun, None)
+
+    def guide():
+        m1 = numpyro.param("m1", 2.0)
+        s1 = numpyro.param("s1", 0.1, constraint=dist.constraints.positive)
+        m2 = numpyro.param("m2", 2.0)
+        s2 = numpyro.param("s2", 0.1, constraint=dist.constraints.positive)
+
+        def true_fun(_):
+            numpyro.sample("x", dist.Normal(m1, s1))
+
+        def false_fun(_):
+            numpyro.sample("x", dist.Normal(m2, s2))
+
+        cluster = numpyro.sample("cluster", dist.Normal())
+        cond(cluster > 0, true_fun, false_fun, None)
+
+    svi = SVI(model, guide, numpyro.optim.Adam(1e-2), Trace_ELBO(num_particles=100))
+    params, losses = svi.run(random.PRNGKey(0), num_steps=2500)
+
+    predictive = Predictive(
+        model,
+        guide=guide,
+        params=params,
+        num_samples=1000,
+        return_sites=["cluster", "x", "z"],
+    )
+    result = predictive(random.PRNGKey(0))
+
+    assert result["cluster"].shape == (1000,)
+    assert result["x"].shape == (1000,)
+    assert result["z"].shape == (1000,)
+
+    mcmc = MCMC(
+        NUTS(model),
+        num_warmup=500,
+        num_samples=2500,
+        num_chains=4,
+        chain_method="sequential",
+    )
+    mcmc.run(random.PRNGKey(0))
+
+    x = mcmc.get_samples()["x"]
+    assert x.shape == (10_000,)
+    assert_allclose(
+        [
+            x.mean(),
+            x.std(),
+            x[x > 2.0].mean(),
+            x[x > 2.0].std(),
+            x[x < 2.0].mean(),
+            x[x < 2.0].std(),
+        ],
+        [2.0, jnp.sqrt(5.0), 4.0, 1.0, 0.0, 1.0],
+        atol=0.1,
+    )
