Add partial evaluation rule for Monoid.plus

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(),
+)

effectful/ops/monoid.py

@@ -892,6 +892,36 @@ 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(is_ready(arg) for arg in args)
+        if not (0 < n_concrete < len(args)):
+            return fwd()
+
+        progress = False
+        new_args = []
+        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:
+                new_args.append(result)
+                progress = True
+        return monoid.plus(*new_args) if progress else fwd()
+
+
 class ReduceFusion(ObjectInterpretation):
     """Implements the identity
     reduce(R, S1, reduce(R, S2, body)) = reduce(R, S1 ∪ S2, body)
@@ -1509,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`."""
 
@@ -2128,6 +2178,7 @@ def extend(self, *intps: Interpretation) -> typing.Self:
     ReduceWeightedStream(),
     ReduceMaskHoist(),
     EliminateSingletonStreams(),
+    PlusPartial(),
     PlusEmpty(),
     PlusSingle(),
     PlusAssoc(),
@@ -2147,6 +2198,16 @@ def extend(self, *intps: Interpretation) -> typing.Self:
     ReduceDependentRangeMask(),
     ReduceDisequalityMask(),
     ContractLongestStream(),
+    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

tests/test_ops_monoid.py

@@ -36,6 +36,7 @@
     PlusEmpty,
     PlusInverseCancellation,
     PlusOrder,
+    PlusPartial,
     PlusSingle,
     Product,
     ReduceDependentRangeMask,
@@ -54,6 +55,7 @@
     ReduceWhereEqualityPeel,
     ReduceWhereToMasks,
     Sum,
+    SumPlus,
     Union,
     WhereHoist,
     as_iterable,
@@ -566,6 +568,25 @@ def test_plus_zero(monoid, backend: Backend):
     backend.check_rewrite(lhs=lhs_left, rhs=rhs, rule={})
 
 
+def test_plus_partial():
+    backend = IntBackend()
+    x = backend.define_vars("x", ret="scalar")
+    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():
+    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")