From aced98a4a035710379279c82e9d04f9d3ff6f91c Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Fri, 7 Aug 2026 16:27:16 -0400 Subject: [PATCH 1/2] add partial evaluation rule for Monoid.plus --- effectful/ops/monoid.py | 29 +++++++++++++++++++++++++++++ tests/test_ops_monoid.py | 29 +++++++++++++++++++++++++++++ 2 files changed, 58 insertions(+) diff --git a/effectful/ops/monoid.py b/effectful/ops/monoid.py index f2cfee75..b1718abb 100644 --- a/effectful/ops/monoid.py +++ b/effectful/ops/monoid.py @@ -892,6 +892,34 @@ def _(self, monoid, body, streams): return fwd() +class PlusPartial(ObjectInterpretation): + @implements(Monoid.plus) + def plus(self, monoid, *args): + """Evaluate maximal concrete runs without reordering symbolic operands.""" + n_concrete = sum(len(fvsof(arg)) == 0 for arg in args) + if n_concrete <= 0 or n_concrete >= len(args): + return fwd() + + progress = False + new_args = [] + run = [] + for arg in args: + if fvsof(arg): + if run: + new_args.append(monoid.plus(*run) if len(run) > 1 else run[0]) + progress |= len(run) > 1 + run = [] + new_args.append(arg) + else: + run.append(arg) + if run: + progress |= len(run) > 1 + new_args.append(monoid.plus(*run) if len(run) > 1 else run[0]) + if not progress: + return fwd() + return monoid.plus(*new_args) + + class ReduceFusion(ObjectInterpretation): """Implements the identity reduce(R, S1, reduce(R, S2, body)) = reduce(R, S1 ∪ S2, body) @@ -2147,6 +2175,7 @@ def extend(self, *intps: Interpretation) -> typing.Self: ReduceDependentRangeMask(), ReduceDisequalityMask(), ContractLongestStream(), + PlusPartial(), ) """``NormalizeIntp`` applies pure-Term rewrites (associativity, distributivity, identity elimination, fusion, factorization, etc.) that drive a reduce diff --git a/tests/test_ops_monoid.py b/tests/test_ops_monoid.py index 21fb8bc7..1ce665fe 100644 --- a/tests/test_ops_monoid.py +++ b/tests/test_ops_monoid.py @@ -36,6 +36,7 @@ PlusEmpty, PlusInverseCancellation, PlusOrder, + PlusPartial, PlusSingle, Product, ReduceDependentRangeMask, @@ -566,6 +567,34 @@ def test_plus_zero(monoid, backend: Backend): backend.check_rewrite(lhs=lhs_left, rhs=rhs, rule={}) +@pytest.mark.parametrize( + ("monoid", "left", "right"), + [ + pytest.param(Sum, 3, 7, id="Sum"), + pytest.param(Product, 2, 12, id="Product"), + pytest.param(Min, 1, 3, id="Min"), + pytest.param(Max, 2, 4, id="Max"), + ], +) +def test_plus_partial(monoid, left, right): + backend = IntBackend() + x = backend.define_vars("x", ret="scalar") + lhs = monoid.plus(1, 2, x(), 3, 4) + rhs = monoid.plus(left, x(), right) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=coproduct(EvaluateIntp, PlusPartial())) + + +def test_plus_partial_without_concrete_rule_is_noop(): + backend = IntBackend() + monoid = Monoid(0, "Custom") + x = backend.define_vars("x", ret="scalar") + term = monoid.plus(1, 2, x(), 3, 4) + + with handler(ReducePartial()): + actual = evaluate(term) + assert syntactic_eq(actual, term) + + @pytest.mark.parametrize("monoid", ALL_MONOIDS) def test_partial_1(monoid, backend: Backend): x = backend.define_vars("x", ret="scalar") From 4158c1fb4e8bcf716cf0c50b724775d2bcd1c25f Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Tue, 11 Aug 2026 12:29:51 -0400 Subject: [PATCH 2/2] evaluate eager arrays too --- effectful/handlers/jax/monoid.py | 27 ++++++++----- effectful/ops/monoid.py | 68 +++++++++++++++++++++++--------- tests/test_ops_monoid.py | 18 +++------ 3 files changed, 73 insertions(+), 40 deletions(-) diff --git a/effectful/handlers/jax/monoid.py b/effectful/handlers/jax/monoid.py index 433b5aee..28233e76 100644 --- a/effectful/handlers/jax/monoid.py +++ b/effectful/handlers/jax/monoid.py @@ -35,6 +35,7 @@ _is_simple_range, complement, is_equality, + is_ready, ) from effectful.ops.monoid import Union as UnionM from effectful.ops.semantics import evaluate, fvsof, fwd, handler, typeof @@ -102,6 +103,12 @@ def _is_jax(t): return fwd() +class ReadyEager(ObjectInterpretation): + @implements(is_ready) + def is_ready(self, expr): + return is_eager_array(expr) or fwd() + + class SumPlusJax(ObjectInterpretation): @implements(Sum.plus) def plus(self, *args): @@ -808,14 +815,6 @@ def einsum( EvaluateIntp.extend( - SumPlusJax(), - SumInverseJax(), - ProductPlusJax(), - MinPlusJax(), - MaxPlusJax(), - LogSumExpPlusJax(), - AndPlusJax(), - OrPlusJax(), IteJax(), MaskJax(), ReduceSumProductContraction(), @@ -825,4 +824,14 @@ def einsum( PlusCastArray(), ) -NormalizeIntp.extend(ReduceArrayGather()) +NormalizeIntp.extend( + ReduceArrayGather(), + ReadyEager(), + SumPlusJax(), + ProductPlusJax(), + MinPlusJax(), + MaxPlusJax(), + LogSumExpPlusJax(), + AndPlusJax(), + OrPlusJax(), +) diff --git a/effectful/ops/monoid.py b/effectful/ops/monoid.py index b1718abb..37ae0bf6 100644 --- a/effectful/ops/monoid.py +++ b/effectful/ops/monoid.py @@ -892,32 +892,34 @@ def _(self, monoid, body, streams): return fwd() +@Operation.define +def is_ready(expr: Expr) -> bool: + return not fvsof(expr) + + class PlusPartial(ObjectInterpretation): @implements(Monoid.plus) def plus(self, monoid, *args): """Evaluate maximal concrete runs without reordering symbolic operands.""" - n_concrete = sum(len(fvsof(arg)) == 0 for arg in args) - if n_concrete <= 0 or n_concrete >= len(args): + n_concrete = sum(is_ready(arg) for arg in args) + if not (0 < n_concrete < len(args)): return fwd() progress = False new_args = [] - run = [] - for arg in args: - if fvsof(arg): - if run: - new_args.append(monoid.plus(*run) if len(run) > 1 else run[0]) - progress |= len(run) > 1 - run = [] - new_args.append(arg) + for ready, group in itertools.groupby(args, key=is_ready): + run = tuple(group) + if not ready or len(run) < 2: + new_args.extend(run) + continue + + result = monoid.plus(*run) + if isinstance(result, Term) and _is_monoid_plus(result.op): + new_args.extend(run) else: - run.append(arg) - if run: - progress |= len(run) > 1 - new_args.append(monoid.plus(*run) if len(run) > 1 else run[0]) - if not progress: - return fwd() - return monoid.plus(*new_args) + new_args.append(result) + progress = True + return monoid.plus(*new_args) if progress else fwd() class ReduceFusion(ObjectInterpretation): @@ -1537,6 +1539,26 @@ def plus(self, *args): return functools.reduce(operator.mul, args) +class AndPlus(ObjectInterpretation): + """Scalar implementation of :data:`And`.""" + + @implements(And.plus) + def plus(self, *args): + if any(isinstance(arg, Term) for arg in args): + return fwd() + return all(args) + + +class OrPlus(ObjectInterpretation): + """Scalar implementation of :data:`Or`.""" + + @implements(Or.plus) + def plus(self, *args): + if any(isinstance(arg, Term) for arg in args): + return fwd() + return any(args) + + class LogSumExpPlus(ObjectInterpretation): """Scalar implementation of :data:`LogSumExp`.""" @@ -2156,6 +2178,7 @@ def extend(self, *intps: Interpretation) -> typing.Self: ReduceWeightedStream(), ReduceMaskHoist(), EliminateSingletonStreams(), + PlusPartial(), PlusEmpty(), PlusSingle(), PlusAssoc(), @@ -2175,7 +2198,16 @@ def extend(self, *intps: Interpretation) -> typing.Self: ReduceDependentRangeMask(), ReduceDisequalityMask(), ContractLongestStream(), - PlusPartial(), + SumPlus(), + MinPlus(), + MaxPlus(), + ProductPlus(), + ArgMinPlus(), + ArgMaxPlus(), + CartesianProductPlus(), + UnionPlus(), + AndPlus(), + OrPlus(), ) """``NormalizeIntp`` applies pure-Term rewrites (associativity, distributivity, identity elimination, fusion, factorization, etc.) that drive a reduce diff --git a/tests/test_ops_monoid.py b/tests/test_ops_monoid.py index 1ce665fe..6bc4b881 100644 --- a/tests/test_ops_monoid.py +++ b/tests/test_ops_monoid.py @@ -55,6 +55,7 @@ ReduceWhereEqualityPeel, ReduceWhereToMasks, Sum, + SumPlus, Union, WhereHoist, as_iterable, @@ -567,21 +568,12 @@ def test_plus_zero(monoid, backend: Backend): backend.check_rewrite(lhs=lhs_left, rhs=rhs, rule={}) -@pytest.mark.parametrize( - ("monoid", "left", "right"), - [ - pytest.param(Sum, 3, 7, id="Sum"), - pytest.param(Product, 2, 12, id="Product"), - pytest.param(Min, 1, 3, id="Min"), - pytest.param(Max, 2, 4, id="Max"), - ], -) -def test_plus_partial(monoid, left, right): +def test_plus_partial(): backend = IntBackend() x = backend.define_vars("x", ret="scalar") - lhs = monoid.plus(1, 2, x(), 3, 4) - rhs = monoid.plus(left, x(), right) - backend.check_rewrite(lhs=lhs, rhs=rhs, rule=coproduct(EvaluateIntp, PlusPartial())) + lhs = Sum.plus(1, 2, x(), 3, 4) + rhs = Sum.plus(3, x(), 7) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=coproduct(PlusPartial(), SumPlus())) def test_plus_partial_without_concrete_rule_is_noop():