Support reduction over distributions with enumerable support

effectful/handlers/jax/monoid.py

@@ -20,6 +20,7 @@
     And,
     CartesianProduct,
     EvaluateIntp,
+    LogSumExp,
     Max,
     Min,
     Monoid,
@@ -33,7 +34,6 @@
     _is_monoid_plus,
     _is_simple_range,
     complement,
-    distributes_over,
     is_equality,
 )
 from effectful.ops.monoid import Union as UnionM
@@ -50,12 +50,6 @@
 logger = logging.getLogger(__name__)
 
 
-LogSumExp = Monoid(name="LogSumExp", identity=jnp.asarray(float("-inf")))
-
-# ``Sum`` in log space is multiplication, which distributes over ``LogSumExp``:
-#   a + logsumexp(b, c) = logsumexp(a + b, a + c)
-distributes_over.register(Sum, LogSumExp)
-
 is_equality.register(jnp.equal)
 for a, b in {
     (jnp.less, jnp.greater),

effectful/handlers/numpyro.py

@@ -14,8 +14,17 @@
 import effectful.handlers.jax.numpy as jnp
 from effectful.handlers.jax import bind_dims, jax_getitem, sizesof, unbind_dims
 from effectful.handlers.jax._handlers import _register_jax_op, is_eager_array
-from effectful.ops.semantics import evaluate, typeof
-from effectful.ops.syntax import defdata, defop
+from effectful.ops.monoid import (
+    LogSumExp,
+    Monoid,
+    NormalizeIntp,
+    Product,
+    Stream,
+    Streams,
+    Sum,
+)
+from effectful.ops.semantics import evaluate, fwd, typeof
+from effectful.ops.syntax import ObjectInterpretation, defdata, deffn, defop, implements
 from effectful.ops.types import NotHandled, Operation, Term
 
 
@@ -259,6 +268,13 @@ def has_rsample(self) -> bool:
             raise NotHandled
         return self._pos_base_dist.has_rsample
 
+    @property
+    @defop
+    def has_enumerate_support(self) -> bool:
+        if not self._is_eager:
+            raise NotHandled
+        return self._pos_base_dist.has_enumerate_support
+
     @property
     @defop
     def event_shape(self) -> tuple[int, ...]:
@@ -387,6 +403,7 @@ def __str__(self):
 batch_shape = _DistributionTerm.batch_shape
 event_shape = _DistributionTerm.event_shape
 has_rsample = _DistributionTerm.has_rsample
+has_enumerate_support = _DistributionTerm.has_enumerate_support
 rsample = _DistributionTerm.rsample
 sample = _DistributionTerm.sample
 log_prob = _DistributionTerm.log_prob
@@ -1175,3 +1192,49 @@ def __init__(self, ty, op, base_dist, reinterpreted_batch_ndims, **kwargs):
 @evaluate.register(dist.Independent)
 def _embed_independent(d: dist.Independent) -> Term[dist.Independent]:
     return Independent(d.base_dist, d.reinterpreted_batch_ndims)
+
+
+@Operation.define
+def distribution_stream(
+    distribution: numpyro.distributions.Distribution,
+) -> Stream[jax.Array]:
+    raise NotHandled
+
+
+class ReduceEnumerableDistribution(ObjectInterpretation):
+    """Distributions with enumerable support turn into weighted reductions of
+    arrays. The weighting used depends on the reduction monoid.
+
+    """
+
+    @implements(Monoid.reduce)
+    def _(self, monoid, body, streams: Streams):
+        for stream_id, stream in streams.items():
+            if not (isinstance(stream, Term) and stream.op == distribution_stream):
+                continue
+
+            dist = stream.args[0]
+            assert isinstance(dist, numpyro.distributions.Distribution)
+            if not dist.has_enumerate_support:
+                continue
+
+            support = dist.enumerate_support(expand=False)
+            value = Operation.define(jax.Array)
+            if monoid == LogSumExp:
+                weighted = Sum.weighted(support, deffn(dist.log_prob(value()), value))
+            elif monoid == Sum:
+                weighted = Product.weighted(
+                    support, deffn(jnp.exp(dist.log_prob(value())), value)
+                )
+            else:
+                continue
+
+            new_streams = {k: v for k, v in streams.items() if k != stream_id} | {
+                stream_id: weighted
+            }
+            return monoid.reduce(body, new_streams)
+
+        return fwd()
+
+
+NormalizeIntp.extend(ReduceEnumerableDistribution())

effectful/ops/monoid.py

@@ -1,6 +1,7 @@
 import collections.abc
 import functools
 import itertools
+import math
 import operator
 import typing
 from collections import UserDict, defaultdict
@@ -183,6 +184,7 @@ def __init__(self, name: str, identity: T, zero: T):
 ArgMax = Monoid(name="ArgMax", identity=(Max.identity, None))
 Sum = Monoid(name="Sum", identity=0)
 Product = MonoidWithZero(name="Product", identity=1, zero=0)
+LogSumExp = Monoid(name="LogSumExp", identity=float("-inf"))
 CartesianProduct: MonoidWithZero[Sequence[Mapping]] = MonoidWithZero(
     name="CartesianProduct", identity=[{}], zero=[]
 )
@@ -257,6 +259,7 @@ def of(self, t: S) -> S | None:
     (Min, Max),
     (Sum, Min),
     (Sum, Max),
+    (Sum, LogSumExp),
     (Product, Sum),
     (CartesianProduct, Union),
     (And, Or),
@@ -1184,6 +1187,20 @@ def plus(self, *args):
         return functools.reduce(operator.mul, args)
 
 
+class LogSumExpPlus(ObjectInterpretation):
+    """Scalar implementation of :data:`LogSumExp`."""
+
+    @implements(LogSumExp.plus)
+    def plus(self, *args):
+        if not _scalar_args(args):
+            return fwd()
+
+        m = max(args)
+        return m + math.log(
+            functools.reduce(operator.add, (math.exp(x - m) for x in args))
+        )
+
+
 class ArgMinPlus(ObjectInterpretation):
     """Scalar score implementation of :data:`ArgMin`."""
 

tests/test_handlers_numpyro.py

@@ -8,12 +8,15 @@
 import numpyro.distributions
 import pytest
 
+import effectful.handlers.jax.monoid  # noqa: F401
 import effectful.handlers.jax.numpy as jnp
 import effectful.handlers.numpyro as dist
 from effectful.handlers.jax import bind_dims, jax_getitem, sizesof, unbind_dims
+from effectful.ops.monoid import LogSumExp, Product, Sum
 from effectful.ops.semantics import typeof
-from effectful.ops.syntax import defop
+from effectful.ops.syntax import deffn, defop
 from effectful.ops.types import Operation, Term
+from tests._monoid_helpers import JaxBackend
 
 ##################################################
 # Test cases
@@ -858,6 +861,41 @@ def test_distribution_support():
     assert isinstance(d.support, numpyro.distributions.constraints.Constraint)
 
 
+@pytest.fixture
+def monoid_backend() -> JaxBackend:
+    return JaxBackend()
+
+
+@pytest.mark.parametrize(
+    ("reduction_monoid", "weight_monoid", "log_weights"),
+    [(Sum, Product, False), (LogSumExp, Sum, True)],
+)
+def test_reduce_enumerable_distribution(
+    reduction_monoid, weight_monoid, log_weights, monoid_backend: JaxBackend
+):
+    value = monoid_backend.define_vars("value", ret="scalar")
+    body = monoid_backend.define_vars(
+        "body", arg_types=(monoid_backend.scalar_typ,), ret="scalar"
+    )
+    distribution = dist.CategoricalProbs(jnp.array([0.25, 0.75]))
+    support = distribution.enumerate_support(expand=False)
+    weight_value = defop(jax.Array, name="weight_value")
+    weight_body = distribution.log_prob(weight_value())
+    if not log_weights:
+        weight_body = jnp.exp(weight_body)
+    weight = deffn(weight_body, weight_value)
+
+    lhs = reduction_monoid.reduce(
+        body(value()), {value: dist.distribution_stream(distribution)}
+    )
+    rhs = reduction_monoid.reduce(
+        body(value()), {value: weight_monoid.weighted(support, weight)}
+    )
+    monoid_backend.check_rewrite(
+        lhs=lhs, rhs=rhs, rule=dist.ReduceEnumerableDistribution()
+    )
+
+
 @pytest.mark.parametrize(
     "dist_factory,dist_args",
     [(dist.Normal, []), (dist.BernoulliProbs, [jnp.array(0.5)])],