diff --git a/effectful/handlers/jax/monoid.py b/effectful/handlers/jax/monoid.py index 433b5aee7..28233e764 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 f2cfee75b..37ae0bf6f 100644 --- a/effectful/ops/monoid.py +++ b/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 diff --git a/tests/test_ops_monoid.py b/tests/test_ops_monoid.py index 21fb8bc7c..6bc4b8810 100644 --- a/tests/test_ops_monoid.py +++ b/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")