From df1f800b5b742fa4695e2014a34fc0f66896ae20 Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Thu, 30 Jul 2026 11:09:59 -0400 Subject: [PATCH 1/6] hoist logsumexp out of the jax module --- effectful/handlers/jax/monoid.py | 6 ------ effectful/ops/monoid.py | 17 +++++++++++++++++ 2 files changed, 17 insertions(+), 6 deletions(-) diff --git a/effectful/handlers/jax/monoid.py b/effectful/handlers/jax/monoid.py index 2e99d6859..ab7212ec2 100644 --- a/effectful/handlers/jax/monoid.py +++ b/effectful/handlers/jax/monoid.py @@ -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), diff --git a/effectful/ops/monoid.py b/effectful/ops/monoid.py index 4830133d3..22dac621c 100644 --- a/effectful/ops/monoid.py +++ b/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`.""" From 1fff442e9f39dd2aef761b616dbfe280aaaaaf17 Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Thu, 30 Jul 2026 11:16:08 -0400 Subject: [PATCH 2/6] add reduction rule for enumerable distributions --- effectful/handlers/numpyro.py | 42 +++++++++++++++++++++++++++++++++-- 1 file changed, 40 insertions(+), 2 deletions(-) diff --git a/effectful/handlers/numpyro.py b/effectful/handlers/numpyro.py index 329cf4719..b0d9691a4 100644 --- a/effectful/handlers/numpyro.py +++ b/effectful/handlers/numpyro.py @@ -14,8 +14,9 @@ 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, Product, Stream, Streams, Sum +from effectful.ops.semantics import evaluate, fwd, typeof +from effectful.ops.syntax import ObjectInterpretation, defdata, defop, implements from effectful.ops.types import NotHandled, Operation, Term @@ -1175,3 +1176,40 @@ 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): + @implements(Monoid.reduce) + def _(self, monoid, body, streams: Streams): + for stream_id, stream in streams.values(): + if not (isinstance(stream, Term) and stream.op == distribution_stream): + continue + + dist = stream.args[0] + assert isinstance(dist, numpyro.distribution.Distribution) + if not dist.has_enumerate_support: + continue + + support = dist.enumerate_support(expand=False) + if monoid == LogSumExp: + weighted = Sum.weighted(support, dist.log_prob) + elif monoid == Sum: + weighted = Product.weighted( + support, lambda x: jnp.exp(dist.log_prob(x)) + ) + 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() From 024319f8fe83ae895e752dc79a108d65546e3a83 Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Thu, 30 Jul 2026 11:30:34 -0400 Subject: [PATCH 3/6] add tests --- effectful/handlers/numpyro.py | 30 +++++++++++++++++++++++++++--- tests/test_handlers_numpyro.py | 32 +++++++++++++++++++++++++++++++- 2 files changed, 58 insertions(+), 4 deletions(-) diff --git a/effectful/handlers/numpyro.py b/effectful/handlers/numpyro.py index b0d9691a4..3eddc4d2b 100644 --- a/effectful/handlers/numpyro.py +++ b/effectful/handlers/numpyro.py @@ -14,7 +14,15 @@ 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.monoid import LogSumExp, Monoid, Product, Stream, Streams, Sum +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, defop, implements from effectful.ops.types import NotHandled, Operation, Term @@ -260,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, ...]: @@ -388,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 @@ -1186,14 +1202,19 @@ def distribution_stream( 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.values(): + 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.distribution.Distribution) + assert isinstance(dist, numpyro.distributions.Distribution) if not dist.has_enumerate_support: continue @@ -1213,3 +1234,6 @@ def _(self, monoid, body, streams: Streams): return monoid.reduce(body, new_streams) return fwd() + + +NormalizeIntp.extend(ReduceEnumerableDistribution()) diff --git a/tests/test_handlers_numpyro.py b/tests/test_handlers_numpyro.py index befd5cbec..f7e3f7e3a 100644 --- a/tests/test_handlers_numpyro.py +++ b/tests/test_handlers_numpyro.py @@ -11,7 +11,8 @@ 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.semantics import typeof +from effectful.ops.monoid import LogSumExp, Product, Sum +from effectful.ops.semantics import evaluate, handler, typeof from effectful.ops.syntax import defop from effectful.ops.types import Operation, Term @@ -858,6 +859,35 @@ def test_distribution_support(): assert isinstance(d.support, numpyro.distributions.constraints.Constraint) +@pytest.mark.parametrize( + ("reduction_monoid", "weight_monoid", "expected_weights"), + [ + (Sum, Product, jnp.array([0.25, 0.75])), + (LogSumExp, Sum, jnp.log(jnp.array([0.25, 0.75]))), + ], +) +def test_reduce_enumerable_distribution( + reduction_monoid, weight_monoid, expected_weights +): + value = defop(jax.Array, name="value") + distribution = dist.CategoricalProbs(jnp.array([0.25, 0.75])) + expression = reduction_monoid.reduce( + value(), {value: dist.distribution_stream(distribution)} + ) + + with handler(dist.ReduceEnumerableDistribution()): + rewritten = evaluate(expression) + + assert isinstance(rewritten, Term) and rewritten.op is reduction_monoid.reduce + rewritten_stream = next(iter(rewritten.args[1].values())) + assert isinstance(rewritten_stream, Term) + assert rewritten_stream.op is weight_monoid.weighted + + support, weight = rewritten_stream.args + assert jnp.array_equal(support, jnp.array([0, 1])) + assert jnp.allclose(weight(support), expected_weights) + + @pytest.mark.parametrize( "dist_factory,dist_args", [(dist.Normal, []), (dist.BernoulliProbs, [jnp.array(0.5)])], From 7bebd8554db41d0dba292acb4e532ee603610037 Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Thu, 30 Jul 2026 11:31:06 -0400 Subject: [PATCH 4/6] lint --- effectful/handlers/jax/monoid.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/effectful/handlers/jax/monoid.py b/effectful/handlers/jax/monoid.py index ab7212ec2..6338e7774 100644 --- a/effectful/handlers/jax/monoid.py +++ b/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 From 991d1df26b30db44c37508a6ccd32008ac985bb1 Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Thu, 30 Jul 2026 11:35:11 -0400 Subject: [PATCH 5/6] rework tests --- effectful/handlers/numpyro.py | 9 ++++-- tests/test_handlers_numpyro.py | 54 +++++++++++++++++++--------------- 2 files changed, 37 insertions(+), 26 deletions(-) diff --git a/effectful/handlers/numpyro.py b/effectful/handlers/numpyro.py index 3eddc4d2b..9155d98be 100644 --- a/effectful/handlers/numpyro.py +++ b/effectful/handlers/numpyro.py @@ -24,7 +24,7 @@ Sum, ) from effectful.ops.semantics import evaluate, fwd, typeof -from effectful.ops.syntax import ObjectInterpretation, defdata, defop, implements +from effectful.ops.syntax import ObjectInterpretation, defdata, deffn, defop, implements from effectful.ops.types import NotHandled, Operation, Term @@ -1219,11 +1219,14 @@ def _(self, monoid, body, streams: Streams): continue support = dist.enumerate_support(expand=False) + value = Operation.define(jax.Array) if monoid == LogSumExp: - weighted = Sum.weighted(support, dist.log_prob) + weighted = Sum.weighted( + support, deffn(dist.log_prob(value()), value) + ) elif monoid == Sum: weighted = Product.weighted( - support, lambda x: jnp.exp(dist.log_prob(x)) + support, deffn(jnp.exp(dist.log_prob(value())), value) ) else: continue diff --git a/tests/test_handlers_numpyro.py b/tests/test_handlers_numpyro.py index f7e3f7e3a..8a9b03b3c 100644 --- a/tests/test_handlers_numpyro.py +++ b/tests/test_handlers_numpyro.py @@ -8,13 +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 evaluate, handler, typeof -from effectful.ops.syntax import defop +from effectful.ops.semantics import typeof +from effectful.ops.syntax import deffn, defop from effectful.ops.types import Operation, Term +from tests._monoid_helpers import JaxBackend ################################################## # Test cases @@ -859,33 +861,39 @@ 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", "expected_weights"), - [ - (Sum, Product, jnp.array([0.25, 0.75])), - (LogSumExp, Sum, jnp.log(jnp.array([0.25, 0.75]))), - ], + ("reduction_monoid", "weight_monoid", "log_weights"), + [(Sum, Product, False), (LogSumExp, Sum, True)], ) def test_reduce_enumerable_distribution( - reduction_monoid, weight_monoid, expected_weights + reduction_monoid, weight_monoid, log_weights, monoid_backend: JaxBackend ): - value = defop(jax.Array, name="value") + 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])) - expression = reduction_monoid.reduce( - value(), {value: dist.distribution_stream(distribution)} + 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() ) - - with handler(dist.ReduceEnumerableDistribution()): - rewritten = evaluate(expression) - - assert isinstance(rewritten, Term) and rewritten.op is reduction_monoid.reduce - rewritten_stream = next(iter(rewritten.args[1].values())) - assert isinstance(rewritten_stream, Term) - assert rewritten_stream.op is weight_monoid.weighted - - support, weight = rewritten_stream.args - assert jnp.array_equal(support, jnp.array([0, 1])) - assert jnp.allclose(weight(support), expected_weights) @pytest.mark.parametrize( From d8e96289c0a6078703423b4467c966d9063cac23 Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Thu, 30 Jul 2026 11:35:57 -0400 Subject: [PATCH 6/6] format --- effectful/handlers/numpyro.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/effectful/handlers/numpyro.py b/effectful/handlers/numpyro.py index 9155d98be..818f5b3bb 100644 --- a/effectful/handlers/numpyro.py +++ b/effectful/handlers/numpyro.py @@ -1221,9 +1221,7 @@ def _(self, monoid, body, streams: Streams): support = dist.enumerate_support(expand=False) value = Operation.define(jax.Array) if monoid == LogSumExp: - weighted = Sum.weighted( - support, deffn(dist.log_prob(value()), value) - ) + 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)