From eadb513d9589d619a30ee87ebfcd12d9f8018985 Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Fri, 7 Aug 2026 13:10:01 -0400 Subject: [PATCH 01/15] add Group and associated rules --- effectful/handlers/jax/monoid.py | 9 +++++++++ effectful/ops/monoid.py | 11 +++++++++++ 2 files changed, 20 insertions(+) diff --git a/effectful/handlers/jax/monoid.py b/effectful/handlers/jax/monoid.py index 5d6d17cf0..433b5aee7 100644 --- a/effectful/handlers/jax/monoid.py +++ b/effectful/handlers/jax/monoid.py @@ -116,6 +116,14 @@ def inverse(self, value): return jnp.negative(value) +class SumInverseJax(ObjectInterpretation): + @implements(Sum.inverse) + def inverse(self, value): + if not _jax_args((value,)): + return fwd() + return jnp.negative(value) + + class ProductPlusJax(ObjectInterpretation): @implements(Product.plus) def plus(self, *args): @@ -801,6 +809,7 @@ def einsum( EvaluateIntp.extend( SumPlusJax(), + SumInverseJax(), ProductPlusJax(), MinPlusJax(), MaxPlusJax(), diff --git a/effectful/ops/monoid.py b/effectful/ops/monoid.py index fec2cdce8..1b6b1d57a 100644 --- a/effectful/ops/monoid.py +++ b/effectful/ops/monoid.py @@ -1271,6 +1271,16 @@ def inverse(self, value): return -value +class SumInverse(ObjectInterpretation): + """Scalar implementation of :meth:`Sum.inverse`.""" + + @implements(Sum.inverse) + def inverse(self, value): + if isinstance(value, Term) or not isinstance(value, int | float): + return fwd() + return -value + + class MinPlus(ObjectInterpretation): """Scalar implementation of :data:`Min`.""" @@ -1842,6 +1852,7 @@ def extend(self, *intps: Interpretation) -> typing.Self: ReducePartial(), DeltaConcrete(), SumPlus(), + SumInverse(), MinPlus(), MaxPlus(), ProductPlus(), From 2edfb8ba21ff74e3feb5ee5b3109e7f34fa49fad Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Tue, 11 Aug 2026 12:02:40 -0400 Subject: [PATCH 02/15] generalize ReduceEqualityMaskRange --- effectful/ops/monoid.py | 195 +++++++++++++++++++++++++++++++-------- tests/test_ops_monoid.py | 156 ++++++++++++++++++++++++------- 2 files changed, 282 insertions(+), 69 deletions(-) diff --git a/effectful/ops/monoid.py b/effectful/ops/monoid.py index 1b6b1d57a..c42a57760 100644 --- a/effectful/ops/monoid.py +++ b/effectful/ops/monoid.py @@ -188,6 +188,12 @@ def __init__(self, name: str, identity: T, zero: T): self.zero = zero +@Operation.define +def as_relation(value: Any) -> Iterable[Any]: + """Cast a domain representation to the common relation carrier type.""" + raise NotHandled + + Min = Monoid(name="Min", identity=float("inf")) Max = Monoid(name="Max", identity=-float("inf")) ArgMin = Monoid(name="ArgMin", identity=(Min.identity, None)) @@ -199,6 +205,11 @@ def __init__(self, name: str, identity: T, zero: T): name="CartesianProduct", identity=[{}], zero=[] ) Union: Monoid[Sequence[Mapping]] = Monoid(name="Union", identity=[]) +Intersection: MonoidWithZero[Iterable[Any]] = MonoidWithZero( + name="Intersection", + identity=Operation.define(Iterable[Any], name="universal")(), + zero=[], +) And = MonoidWithZero(name="And", identity=True, zero=False) Or = Monoid(name="Or", identity=False) @@ -228,7 +239,7 @@ def __call__(self, t: T) -> bool: is_commutative = _ExtensiblePredicate({Max, Min, Sum, Product, And, Or}) -is_idempotent = _ExtensiblePredicate({Max, Min, And, Or}) +is_idempotent = _ExtensiblePredicate({Max, Min, And, Or, Intersection}) @dataclass @@ -273,6 +284,7 @@ def of(self, t: S) -> S | None: (Product, Sum), (CartesianProduct, Union), (And, Or), + (Intersection, Union), ) @@ -391,66 +403,74 @@ def group_plus(value): class ReduceEqualityMaskRange(ObjectInterpretation): - """M.reduce(M.mask(v, And.plus(i = x, *m)), {i: range(N)} ∪ S) ≡ - M.mask(M.reduce(M.mask(v, *m), {i: [x]} ∪ S), And.plus(0 <= x, x < N)) + """Eliminate an equality by intersecting a stream with a singleton. - The equality constraint ``i = x`` on a range-stream reduce is discharged by - a gather (the stream becomes the singleton ``[x]``) guarded by a bounds - check. + For either orientation of an equality between a reduced variable and an + expression that is not another reduced-variable symbol:: + + M.reduce(M.mask(v, And.plus(i == x, *m)), {i: I} | S) + == M.reduce(M.mask(v, And.plus(*m)), + {i: Intersection.plus(I, [x])} | S) + + The expression ``x`` may depend on streams in ``S``. Keeping those streams + in the outer bundle makes the intersection pointwise, so this rewrite does + not require idempotence or source-variable liveness checks. A separate + :class:`ReduceIntersectionSingletonRange` rule lowers intersections with + simple ranges to singleton gathers guarded by bounds masks. When the reduce body is a ``plus`` of the same monoid, the rule distributes the reduce over the plus -- but only when doing so exposes an eliminable equality mask in some summand. This is a *targeted* split (it leaves ``ReduceSplit`` conservative): summands whose reduced index appears in an - equality become gathers, while the rest stay as ordinary masked reduces. + equality get restricted domains, while the rest stay as ordinary masked + reduces. """ @staticmethod def _match_eq(cond, streams): - """If ``cond`` is ``stream_op == key`` (either order) where ``stream_op`` - is a ``range(0, N)`` stream and ``key`` is stream-independent, return - ``(stream_op, key)``; otherwise ``None``.""" + """Match ``stream_op == expr`` in either orientation. + + Equalities between two bare reduced-variable symbols are left alone; + choosing which domain should represent their intersection requires a + separate rule. + """ - def test(op, stream_op, mask_key): + def is_stream_symbol(term): return ( - is_equality(op) - and stream_op in streams - and _is_simple_range(streams[stream_op]) - and not (fvsof(mask_key) & set(streams)) + isinstance(term, Term) + and not term.args + and not term.kwargs + and term.op in streams ) - match cond: - case Term(op, (Term(stream_op, (), {}), mask_key), {}) if test( - op, stream_op, mask_key - ): - return (stream_op, mask_key) - case Term(op, (mask_key, Term(stream_op, (), {})), {}) if test( - op, stream_op, mask_key - ): - return (stream_op, mask_key) - case _: - return None + if not (isinstance(cond, Term) and is_equality(cond.op)): + return None + + lhs, rhs = cond.args + for stream_term, expr in ((lhs, rhs), (rhs, lhs)): + if is_stream_symbol(stream_term) and not is_stream_symbol(expr): + return stream_term.op, expr + return None def _eliminate(self, monoid, value, mask, streams): - """Discharge one eliminable equality constraint via a gather, or return - ``None`` if no constraint is eliminable.""" + """Discharge one eliminable equality constraint, or return ``None``.""" conds = _conjuncts(mask) for i, cond in enumerate(conds): matched = self._match_eq(cond, streams) if matched is None: continue - stream_op, mask_key = matched - stream = streams[stream_op] + stream_op, expr = matched return monoid.reduce( monoid.mask( - monoid.mask( - value, - And.plus(stream.start <= mask_key, mask_key < stream.stop), - ), + value, And.plus(*(c for (j, c) in enumerate(conds) if i != j)), ), - {stream_op: (mask_key,)} - | {k: v for (k, v) in streams.items() if k != stream_op}, + { + stream_op: Intersection.plus( + as_relation(streams[stream_op]), as_relation([expr]) + ), + } + | {k: v for (k, v) in streams.items() if k is not stream_op}, ) return None @@ -482,6 +502,54 @@ def _(self, monoid, body, streams): return fwd() +class ReduceIntersectionSingletonRange(ObjectInterpretation): + """Lower the intersection of a simple range and a singleton stream. + + M.reduce(v, {i: Intersection.plus(range(N), [x])} | S) + == M.reduce(M.mask(v, 0 <= x < N), {i: (x,)} | S) + + The singleton expression may depend on streams in ``S``; the resulting + dependent singleton is eliminated later by :class:`EliminateSingletonStreams`. + """ + + @staticmethod + def _match(stream): + match stream: + case Term(op, (lhs, rhs), {}) if op is Intersection.plus: + pass + case _: + return None + + def uncast(arg): + if isinstance(arg, Term) and arg.op is as_relation: + return arg.args[0] + return arg + + lhs, rhs = uncast(lhs), uncast(rhs) + if _is_simple_range(lhs) and isinstance(rhs, list) and len(rhs) == 1: + return lhs, rhs[0] + if _is_simple_range(rhs) and isinstance(lhs, list) and len(lhs) == 1: + return rhs, lhs[0] + return None + + @implements(Monoid.reduce) + def reduce(self, monoid, body, streams): + for stream_op, stream in streams.items(): + matched = self._match(stream) + if matched is None: + continue + range_stream, value = matched + return monoid.reduce( + monoid.mask( + body, + And.plus(range_stream.start <= value, value < range_stream.stop), + ), + {stream_op: (value,)} + | {k: v for k, v in streams.items() if k is not stream_op}, + ) + return fwd() + + class ReduceMaskHoist(ObjectInterpretation): """M.reduce(M.mask(v, c), S) ≡ M.mask(M.reduce(v, S), c) when ``c`` does not depend on any stream in ``S``. @@ -1391,6 +1459,58 @@ def plus(self, *args): return list(itertools.chain(*args)) +class PlusCastIntersection(ObjectInterpretation): + """Cast heterogeneous intersection arguments to a common relation type.""" + + @implements(Intersection.plus) + def plus(self, *args): + typs = [typeof(arg) for arg in args] + if not args or all(typ == typs[0] for typ in typs[1:]): + return fwd() + return Intersection.plus(*(defdata(as_relation, arg) for arg in args)) + + +class IntersectionPlus(ObjectInterpretation): + """Pure-Python filtering implementation of :data:`Intersection`. + + This preserves occurrences from the leftmost stream. Array-valued elements + remain symbolic so backend-specific, pointwise intersection lowering can + handle them. + """ + + @staticmethod + def _unwrap(arg): + if isinstance(arg, Term) and arg.op is as_relation: + return arg.args[0] + return arg + + @staticmethod + def _concrete_value(value): + if isinstance(value, tuple): + return all(IntersectionPlus._concrete_value(v) for v in value) + return isinstance(value, bool | int | float | complex | str | bytes) + + @implements(Intersection.plus) + def plus(self, *args): + args = tuple(self._unwrap(arg) for arg in args) + if not args or any(isinstance(arg, Term) for arg in args): + return fwd() + if not all(isinstance(arg, Iterable) for arg in args): + return fwd() + + values = list(args[0]) + tails = [list(arg) for arg in args[1:]] + if not all( + self._concrete_value(value) for value in itertools.chain(values, *tails) + ): + return fwd() + return [ + value + for value in values + if all(any(syntactic_eq(value, other) for other in tail) for tail in tails) + ] + + is_scalar = _ExtensiblePredicate({Min, Max, Sum, Product, And, Or}) @@ -1860,7 +1980,10 @@ def extend(self, *intps: Interpretation) -> typing.Self: ArgMaxPlus(), CartesianProductPlus(), UnionPlus(), + PlusCastIntersection(), + IntersectionPlus(), ReduceEqualityMaskRange(), + ReduceIntersectionSingletonRange(), ReduceWhereToMasks(), ) diff --git a/tests/test_ops_monoid.py b/tests/test_ops_monoid.py index b57226a93..8b246e926 100644 --- a/tests/test_ops_monoid.py +++ b/tests/test_ops_monoid.py @@ -18,6 +18,7 @@ EvaluateIntp, Factor, Group, + Intersection, InverseInverse, InversePlus, Max, @@ -41,6 +42,7 @@ ReduceEmpty, ReduceEqualityMaskRange, ReduceFusion, + ReduceIntersectionSingletonRange, ReduceMaskHoist, ReducePartial, ReduceSplit, @@ -52,6 +54,7 @@ Sum, Union, WhereHoist, + as_relation, distributes_over, is_commutative, solve_group_equality, @@ -741,28 +744,21 @@ def test_reduce_mask_hoist_dependent_noop(monoid): @pytest.mark.parametrize("monoid", ALL_MONOIDS) def test_reduce_equality_mask_range_simple(backend: Backend, monoid): - """Best case: a single equality on the reduced range stream becomes a - singleton-stream gather guarded by the corresponding bounds check. - """ + """An equality restricts its stream to an intersection with a singleton.""" a, c = backend.define_vars("a", "c", ret="scalar") f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="scalar") lhs = monoid.reduce(monoid.mask(f(a()), a() == c()), {a: range(3)}) rhs = monoid.reduce( - monoid.mask( - monoid.mask(f(a()), And.plus(0 <= c(), c() < 3)), - And.plus(), - ), - {a: (c(),)}, + monoid.mask(f(a()), And.plus()), + {a: Intersection.plus(as_relation(range(3)), as_relation([c()]))}, ) backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceEqualityMaskRange()) @pytest.mark.parametrize("monoid", ALL_MONOIDS) def test_reduce_equality_mask_range_residual_conjuncts(backend: Backend, monoid): - """Non-equality conjuncts are preserved inside the gathered singleton - reduce; only the equality on the reduced range stream is discharged. - """ + """Only the equality is moved into the stream-domain intersection.""" a, c, d, e = backend.define_vars("a", "c", "d", "e", ret="scalar") f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="scalar") @@ -771,48 +767,142 @@ def test_reduce_equality_mask_range_residual_conjuncts(backend: Backend, monoid) {a: range(4)}, ) rhs = monoid.reduce( - monoid.mask( - monoid.mask(f(a()), And.plus(0 <= c(), c() < 4)), - And.plus(d() < e(), c() < e()), - ), - {a: (c(),)}, + monoid.mask(f(a()), And.plus(d() < e(), c() < e())), + {a: Intersection.plus(as_relation(range(4)), as_relation([c()]))}, ) backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceEqualityMaskRange()) +def test_reduce_equality_mask_dependent_intersection(): + """The singleton expression may depend on another retained stream.""" + backend = IntBackend() + x, y, c, d = backend.define_vars("x", "y", "c", "d", ret="scalar") + X, Y = backend.define_vars("X", "Y", ret="stream") + f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="scalar") + g = backend.define_vars("g", arg_types=(backend.scalar_typ,), ret="scalar") + + lhs = Min.reduce( + Min.mask(g(x()), And.plus(x() == f(y()), c() < d())), + {x: X(), y: Y()}, + ) + rhs = Min.reduce( + Min.mask(g(x()), And.plus(c() < d())), + { + x: Intersection.plus(as_relation(X()), as_relation([f(y())])), + y: Y(), + }, + ) + + with handler(ReduceEqualityMaskRange()): + actual = evaluate(lhs) + assert syntactic_eq_alpha(actual, rhs) + + +def test_reduce_equality_mask_dependent_intersection_nonidempotent(): + """Pointwise intersection does not require an idempotent outer monoid.""" + backend = IntBackend() + x, y = backend.define_vars("x", "y", ret="scalar") + X, Y = backend.define_vars("X", "Y", ret="stream") + f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="scalar") + g = backend.define_vars("g", arg_types=(backend.scalar_typ,), ret="scalar") + + lhs = Sum.reduce(Sum.mask(g(x()), x() == f(y())), {x: X(), y: Y()}) + rhs = Sum.reduce( + Sum.mask(g(x()), And.plus()), + { + x: Intersection.plus(as_relation(X()), as_relation([f(y())])), + y: Y(), + }, + ) + with handler(ReduceEqualityMaskRange()): + actual = evaluate(lhs) + assert syntactic_eq_alpha(actual, rhs) + + +def test_reduce_equality_mask_dependent_intersection_retains_source_uses(): + """Source variables remain available to the body after the rewrite.""" + backend = IntBackend() + x, y = backend.define_vars("x", "y", ret="scalar") + X, Y = backend.define_vars("X", "Y", ret="stream") + f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="scalar") + g = backend.define_vars( + "g", arg_types=(backend.scalar_typ, backend.scalar_typ), ret="scalar" + ) + + lhs = Min.reduce(Min.mask(g(x(), y()), x() == f(y())), {x: X(), y: Y()}) + rhs = Min.reduce( + Min.mask(g(x(), y()), And.plus()), + { + x: Intersection.plus(as_relation(X()), as_relation([f(y())])), + y: Y(), + }, + ) + with handler(ReduceEqualityMaskRange()): + actual = evaluate(lhs) + assert syntactic_eq_alpha(actual, rhs) + + +def test_reduce_equality_mask_image_domain_symbol_side_noop(): + """Bare stream equalities require a separate domain-intersection rule.""" + backend = IntBackend() + x, y = backend.define_vars("x", "y", ret="scalar") + X, Y = backend.define_vars("X", "Y", ret="stream") + g = backend.define_vars("g", arg_types=(backend.scalar_typ,), ret="scalar") + + term = Min.reduce(Min.mask(g(x()), x() == y()), {x: X(), y: Y()}) + with handler(ReduceEqualityMaskRange()): + actual = evaluate(term) + assert syntactic_eq_alpha(actual, term) + + @pytest.mark.parametrize("monoid", ALL_MONOIDS) -def test_reduce_equality_mask_range_noncanonical_range_noop(backend: Backend, monoid): - """The rule only handles ``range(0, N, 1)`` streams; for other ranges it - should leave the term unchanged. - """ +def test_reduce_equality_mask_noncanonical_range(backend: Backend, monoid): + """Equality elimination itself is independent of the domain representation.""" a, c = backend.define_vars("a", "c", ret="scalar") f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="scalar") - term = monoid.reduce(monoid.mask(f(a()), a() == c()), {a: range(1, 4)}) - backend.check_rewrite(lhs=term, rhs=term, rule=ReduceEqualityMaskRange()) + lhs = monoid.reduce(monoid.mask(f(a()), a() == c()), {a: range(1, 4)}) + rhs = monoid.reduce( + monoid.mask(f(a()), And.plus()), + {a: Intersection.plus(as_relation(range(1, 4)), as_relation([c()]))}, + ) + with handler(ReduceEqualityMaskRange()): + actual = evaluate(lhs) + assert syntactic_eq_alpha(actual, rhs) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_reduce_intersection_singleton_range(backend: Backend, monoid): + """A simple-range intersection lowers to a bounds-guarded singleton.""" + x, y = backend.define_vars("x", "y", ret="scalar") + f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="scalar") + + lhs = monoid.reduce( + f(x()), + {x: Intersection.plus(as_relation(range(3)), as_relation([y()]))}, + ) + rhs = monoid.reduce( + monoid.mask(f(x()), And.plus(0 <= y(), y() < 3)), + {x: (y(),)}, + ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceIntersectionSingletonRange()) @pytest.mark.parametrize("monoid", ALL_MONOIDS) def test_reduce_equality_mask_plus(backend: Backend, monoid): - """ReduceEqualityMaskRange distributes over a plus body, discharging an - equality on the reduced stream in one summand via a singleton-stream gather - while leaving the other summand as an ordinary masked reduce. - """ + """Targeted splitting exposes an intersection in the matching summand.""" a, c = backend.define_vars("a", "c", ret="scalar") f, g = backend.define_vars("f", "g", arg_types=(backend.scalar_typ,), ret="scalar") body = monoid.plus( - monoid.mask(f(a()), a() == c()), # eliminable: a == c over range - monoid.mask(g(a()), c() == 0), # not eliminable (no reduced-stream eq) + monoid.mask(f(a()), a() == c()), + monoid.mask(g(a()), c() == 0), ) lhs = monoid.reduce(body, {a: range(3)}) rhs = monoid.plus( monoid.reduce( - monoid.mask( - monoid.mask(f(a()), And.plus(0 <= c(), c() < 3)), - And.plus(), - ), - {a: (c(),)}, + monoid.mask(f(a()), And.plus()), + {a: Intersection.plus(as_relation(range(3)), as_relation([c()]))}, ), monoid.reduce(monoid.mask(g(a()), c() == 0), {a: range(3)}), ) From ae4d13d5a1826af9a598e15ae27b43d8b2c6ac45 Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Mon, 3 Aug 2026 15:48:58 -0400 Subject: [PATCH 03/15] cleanup --- effectful/ops/monoid.py | 119 ++++++++++++++++----------------------- tests/_monoid_helpers.py | 44 +++++++-------- tests/test_ops_monoid.py | 105 +++++++++++++++++----------------- 3 files changed, 118 insertions(+), 150 deletions(-) diff --git a/effectful/ops/monoid.py b/effectful/ops/monoid.py index c42a57760..1298d4bfd 100644 --- a/effectful/ops/monoid.py +++ b/effectful/ops/monoid.py @@ -5,7 +5,14 @@ import operator import typing from collections import UserDict, defaultdict -from collections.abc import Callable, Generator, Iterable, Mapping, Sequence, Sized +from collections.abc import ( + Callable, + Generator, + Iterable, + Mapping, + Sequence, + Sized, +) from dataclasses import dataclass from graphlib import TopologicalSorter from typing import Annotated, Any @@ -188,12 +195,6 @@ def __init__(self, name: str, identity: T, zero: T): self.zero = zero -@Operation.define -def as_relation(value: Any) -> Iterable[Any]: - """Cast a domain representation to the common relation carrier type.""" - raise NotHandled - - Min = Monoid(name="Min", identity=float("inf")) Max = Monoid(name="Max", identity=-float("inf")) ArgMin = Monoid(name="ArgMin", identity=(Min.identity, None)) @@ -204,16 +205,27 @@ def as_relation(value: Any) -> Iterable[Any]: CartesianProduct: MonoidWithZero[Sequence[Mapping]] = MonoidWithZero( name="CartesianProduct", identity=[{}], zero=[] ) -Union: Monoid[Sequence[Mapping]] = Monoid(name="Union", identity=[]) -Intersection: MonoidWithZero[Iterable[Any]] = MonoidWithZero( +Union: Monoid[Iterable] = Monoid(name="Union", identity=[]) +Intersection: MonoidWithZero[Iterable] = MonoidWithZero( name="Intersection", - identity=Operation.define(Iterable[Any], name="universal")(), + identity=Operation.define(Iterable, name="universal")(), zero=[], ) And = MonoidWithZero(name="And", identity=True, zero=False) Or = Monoid(name="Or", identity=False) +@Operation.define +def as_iterable[T](value: Iterable[T]) -> Iterable[T]: + raise NotHandled + + +def _unwrap_as_iterable[T](arg: Iterable[T]) -> Iterable[T]: + if isinstance(arg, Term) and arg.op is as_iterable: + return arg.args[0] # type: ignore + return arg + + def _conjuncts(mask) -> Sequence[Term]: """Return the conjuncts of an ``And`` mask as a flat tuple.""" match mask: @@ -460,18 +472,15 @@ def _eliminate(self, monoid, value, mask, streams): if matched is None: continue stream_op, expr = matched - return monoid.reduce( - monoid.mask( - value, - And.plus(*(c for (j, c) in enumerate(conds) if i != j)), - ), - { - stream_op: Intersection.plus( - as_relation(streams[stream_op]), as_relation([expr]) - ), - } - | {k: v for (k, v) in streams.items() if k is not stream_op}, + + new_mask = monoid.mask( + value, + And.plus(*(c for (j, c) in enumerate(conds) if i != j)), ) + new_streams = { + stream_op: Intersection.plus(streams[stream_op], [expr]), + } | {k: v for (k, v) in streams.items() if k is not stream_op} + return monoid.reduce(new_mask, new_streams) return None def _summand_eliminable(self, monoid, summand, streams): @@ -520,12 +529,7 @@ def _match(stream): case _: return None - def uncast(arg): - if isinstance(arg, Term) and arg.op is as_relation: - return arg.args[0] - return arg - - lhs, rhs = uncast(lhs), uncast(rhs) + lhs, rhs = _unwrap_as_iterable(lhs), _unwrap_as_iterable(rhs) if _is_simple_range(lhs) and isinstance(rhs, list) and len(rhs) == 1: return lhs, rhs[0] if _is_simple_range(rhs) and isinstance(lhs, list) and len(lhs) == 1: @@ -1447,30 +1451,21 @@ def plus(self, *args): return [_disjoint_merge(*vals) for vals in itertools.product(*args)] -class UnionPlus(ObjectInterpretation): - @implements(Union.plus) - def plus(self, *args): - if not args: - return fwd() - if any(isinstance(x, Term) for x in args): - return fwd() - if not all(isinstance(x, Iterable) for x in args): - return fwd() - return list(itertools.chain(*args)) +class PlusCastIterable(ObjectInterpretation): + """Cast heterogeneous iterable arguments to a common iterable type.""" + @implements(Monoid.plus) + def _plus(self, monoid, *args): + if monoid not in (Intersection, Union): + return fwd() -class PlusCastIntersection(ObjectInterpretation): - """Cast heterogeneous intersection arguments to a common relation type.""" - - @implements(Intersection.plus) - def plus(self, *args): - typs = [typeof(arg) for arg in args] - if not args or all(typ == typs[0] for typ in typs[1:]): + if not args or all(isinstance(a, Term) and a.op == as_iterable for a in args): return fwd() - return Intersection.plus(*(defdata(as_relation, arg) for arg in args)) + return monoid.plus(*(as_iterable(arg) for arg in args)) -class IntersectionPlus(ObjectInterpretation): + +class IterablePlus(ObjectInterpretation): """Pure-Python filtering implementation of :data:`Intersection`. This preserves occurrences from the leftmost stream. Array-valued elements @@ -1478,24 +1473,10 @@ class IntersectionPlus(ObjectInterpretation): handle them. """ - @staticmethod - def _unwrap(arg): - if isinstance(arg, Term) and arg.op is as_relation: - return arg.args[0] - return arg - - @staticmethod - def _concrete_value(value): - if isinstance(value, tuple): - return all(IntersectionPlus._concrete_value(v) for v in value) - return isinstance(value, bool | int | float | complex | str | bytes) - @implements(Intersection.plus) - def plus(self, *args): - args = tuple(self._unwrap(arg) for arg in args) - if not args or any(isinstance(arg, Term) for arg in args): - return fwd() - if not all(isinstance(arg, Iterable) for arg in args): + def _intersection_plus(self, *args): + args = tuple(_unwrap_as_iterable(arg) for arg in args) + if not args or any(isinstance(x, Term) for x in args): return fwd() values = list(args[0]) @@ -1504,11 +1485,8 @@ def plus(self, *args): self._concrete_value(value) for value in itertools.chain(values, *tails) ): return fwd() - return [ - value - for value in values - if all(any(syntactic_eq(value, other) for other in tail) for tail in tails) - ] + + return list(itertools.chain(*args)) is_scalar = _ExtensiblePredicate({Min, Max, Sum, Product, And, Or}) @@ -1979,9 +1957,7 @@ def extend(self, *intps: Interpretation) -> typing.Self: ArgMinPlus(), ArgMaxPlus(), CartesianProductPlus(), - UnionPlus(), - PlusCastIntersection(), - IntersectionPlus(), + IterablePlus(), ReduceEqualityMaskRange(), ReduceIntersectionSingletonRange(), ReduceWhereToMasks(), @@ -2031,6 +2007,7 @@ def extend(self, *intps: Interpretation) -> typing.Self: PlusConsecutiveDups(), PlusOrder(), PlusCastFloat(), + PlusCastIterable(), MaskFusion(), MaskBool(), WhereHoist(), diff --git a/tests/_monoid_helpers.py b/tests/_monoid_helpers.py index 25490411e..91f74cb13 100644 --- a/tests/_monoid_helpers.py +++ b/tests/_monoid_helpers.py @@ -168,19 +168,11 @@ def __init__(self): def eq(self, a: Any, b: Any) -> bool: raise NotImplementedError - @abstractmethod - def strategy( - self, - arg_types: tuple[type, ...] = (), - ret: Literal["scalar", "stream"] = "scalar", - ) -> SearchStrategy: - raise NotImplementedError - def _fresh_op( self, name: str, arg_types: tuple[type, ...] = (), - ret: Literal["scalar", "stream"] = "scalar", + ret: Literal["scalar", "stream"] | type = "scalar", ) -> Operation: """Build a fresh, unhandled Operation whose parameter and return annotations are derived from this backend. @@ -190,7 +182,7 @@ def _fresh_op( each of type ``scalar_typ``. """ scalar = self.scalar_typ - out = self.stream_typ if ret == "stream" else scalar + out = self.stream_typ if ret == "stream" else scalar if ret == "scalar" else ret params = ", ".join(f"_a{i}" for i in range(len(arg_types))) ns: dict[str, Any] = {"NotHandled": NotHandled} exec(f"def _fn({params}):\n raise NotHandled\n", ns) @@ -328,28 +320,28 @@ class IntBackend(Backend): lambda x: [0, x, x + 1], ] + @staticmethod def strategy( - self, arg_types: tuple[type, ...] = (), - ret: Literal["scalar", "stream"] = "scalar", + ret: Literal["scalar", "stream"] | type = "scalar", ) -> SearchStrategy: match arg_types, ret: - case (), "scalar": + case (), ("scalar" | builtins.int): return st.integers(min_value=-100, max_value=100).map(deffn) case (), "stream": scalars = st.integers(min_value=-100, max_value=100) return st.lists(scalars, max_size=2).map(deffn) case (builtins.int,), "scalar": - return st.sampled_from(self._unary_num_fns) + return st.sampled_from(IntBackend._unary_num_fns) case (builtins.int, builtins.int), "scalar": - return st.sampled_from(self._binary_num_fns) + return st.sampled_from(IntBackend._binary_num_fns) case (builtins.int, builtins.int, builtins.int), "scalar": return st.tuples( - st.sampled_from(self._binary_num_fns), - st.sampled_from(self._binary_num_fns), + st.sampled_from(IntBackend._binary_num_fns), + st.sampled_from(IntBackend._binary_num_fns), ).map(lambda fg: lambda a, b, c: fg[0](a, fg[1](b, c))) case (builtins.int,), "stream": - return st.sampled_from(self._unary_list_fns) + return st.sampled_from(IntBackend._unary_list_fns) raise NotImplementedError( f"No int strategy for op with return {ret!r} and {arg_types} args" ) @@ -385,12 +377,14 @@ class JaxBackend(Backend): lambda a, b: a * b, ] + @staticmethod def strategy( - self, arg_types: tuple[type, ...] = (), - ret: Literal["scalar", "stream"] = "scalar", + ret: Literal["scalar", "stream"] | type = "scalar", ) -> st.SearchStrategy[Callable]: match arg_types, ret: + case (), builtins.int: + return IntBackend.strategy((), builtins.int) case (), "scalar": return ( st.lists( @@ -412,16 +406,16 @@ def strategy( .map(deffn) ) case (jax.Array,), "scalar": - return st.sampled_from(self._unary_jax_scalar_fns) + return st.sampled_from(JaxBackend._unary_jax_scalar_fns) case (jax.Array, jax.Array), "scalar": - return st.sampled_from(self._binary_jax_scalar_fns) + return st.sampled_from(JaxBackend._binary_jax_scalar_fns) case (jax.Array, jax.Array, jax.Array), "scalar": return st.tuples( - st.sampled_from(self._binary_jax_scalar_fns), - st.sampled_from(self._binary_jax_scalar_fns), + st.sampled_from(JaxBackend._binary_jax_scalar_fns), + st.sampled_from(JaxBackend._binary_jax_scalar_fns), ).map(lambda fg: lambda a, b, c: fg[0](a, fg[1](b, c))) case (jax.Array,), "stream": - return st.sampled_from(self._unary_jax_stream_fns) + return st.sampled_from(JaxBackend._unary_jax_stream_fns) raise NotImplementedError( f"No jax strategy for op with return {ret!r} and {arg_types} args" diff --git a/tests/test_ops_monoid.py b/tests/test_ops_monoid.py index 8b246e926..b83cb11ce 100644 --- a/tests/test_ops_monoid.py +++ b/tests/test_ops_monoid.py @@ -29,6 +29,7 @@ NormalizeIntp, Or, PlusAssoc, + PlusCastIterable, PlusConsecutiveDups, PlusDistr, PlusEmpty, @@ -54,7 +55,7 @@ Sum, Union, WhereHoist, - as_relation, + as_iterable, distributes_over, is_commutative, solve_group_equality, @@ -745,32 +746,34 @@ def test_reduce_mask_hoist_dependent_noop(monoid): @pytest.mark.parametrize("monoid", ALL_MONOIDS) def test_reduce_equality_mask_range_simple(backend: Backend, monoid): """An equality restricts its stream to an intersection with a singleton.""" - a, c = backend.define_vars("a", "c", ret="scalar") - f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="scalar") + a, c = backend.define_vars("a", "c", ret=int) - lhs = monoid.reduce(monoid.mask(f(a()), a() == c()), {a: range(3)}) + lhs = monoid.reduce(monoid.mask(a(), a() == c()), {a: range(3)}) rhs = monoid.reduce( - monoid.mask(f(a()), And.plus()), - {a: Intersection.plus(as_relation(range(3)), as_relation([c()]))}, + monoid.mask(a(), And.plus()), + {a: Intersection.plus(as_iterable(range(3)), as_iterable([c()]))}, + ) + backend.check_rewrite( + lhs=lhs, rhs=rhs, rule=coproduct(ReduceEqualityMaskRange(), PlusCastIterable()) ) - backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceEqualityMaskRange()) @pytest.mark.parametrize("monoid", ALL_MONOIDS) def test_reduce_equality_mask_range_residual_conjuncts(backend: Backend, monoid): """Only the equality is moved into the stream-domain intersection.""" - a, c, d, e = backend.define_vars("a", "c", "d", "e", ret="scalar") - f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="scalar") + a, c, d, e = backend.define_vars("a", "c", "d", "e", ret=int) lhs = monoid.reduce( - monoid.mask(f(a()), And.plus(d() < e(), a() == c(), c() < e())), + monoid.mask(a(), And.plus(d() < e(), a() == c(), c() < e())), {a: range(4)}, ) rhs = monoid.reduce( - monoid.mask(f(a()), And.plus(d() < e(), c() < e())), - {a: Intersection.plus(as_relation(range(4)), as_relation([c()]))}, + monoid.mask(a(), And.plus(d() < e(), c() < e())), + {a: Intersection.plus(as_iterable(range(4)), as_iterable([c()]))}, + ) + backend.check_rewrite( + lhs=lhs, rhs=rhs, rule=coproduct(ReduceEqualityMaskRange(), PlusCastIterable()) ) - backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceEqualityMaskRange()) def test_reduce_equality_mask_dependent_intersection(): @@ -788,14 +791,14 @@ def test_reduce_equality_mask_dependent_intersection(): rhs = Min.reduce( Min.mask(g(x()), And.plus(c() < d())), { - x: Intersection.plus(as_relation(X()), as_relation([f(y())])), + x: Intersection.plus(as_iterable(X()), as_iterable([f(y())])), y: Y(), }, ) - with handler(ReduceEqualityMaskRange()): - actual = evaluate(lhs) - assert syntactic_eq_alpha(actual, rhs) + backend.check_rewrite( + lhs=lhs, rhs=rhs, rule=coproduct(ReduceEqualityMaskRange(), PlusCastIterable()) + ) def test_reduce_equality_mask_dependent_intersection_nonidempotent(): @@ -809,14 +812,11 @@ def test_reduce_equality_mask_dependent_intersection_nonidempotent(): lhs = Sum.reduce(Sum.mask(g(x()), x() == f(y())), {x: X(), y: Y()}) rhs = Sum.reduce( Sum.mask(g(x()), And.plus()), - { - x: Intersection.plus(as_relation(X()), as_relation([f(y())])), - y: Y(), - }, + {x: Intersection.plus(as_iterable(X()), as_iterable([f(y())])), y: Y()}, + ) + backend.check_rewrite( + lhs=lhs, rhs=rhs, rule=coproduct(ReduceEqualityMaskRange(), PlusCastIterable()) ) - with handler(ReduceEqualityMaskRange()): - actual = evaluate(lhs) - assert syntactic_eq_alpha(actual, rhs) def test_reduce_equality_mask_dependent_intersection_retains_source_uses(): @@ -832,14 +832,11 @@ def test_reduce_equality_mask_dependent_intersection_retains_source_uses(): lhs = Min.reduce(Min.mask(g(x(), y()), x() == f(y())), {x: X(), y: Y()}) rhs = Min.reduce( Min.mask(g(x(), y()), And.plus()), - { - x: Intersection.plus(as_relation(X()), as_relation([f(y())])), - y: Y(), - }, + {x: Intersection.plus(as_iterable(X()), as_iterable([f(y())])), y: Y()}, + ) + backend.check_rewrite( + lhs=lhs, rhs=rhs, rule=coproduct(ReduceEqualityMaskRange(), PlusCastIterable()) ) - with handler(ReduceEqualityMaskRange()): - actual = evaluate(lhs) - assert syntactic_eq_alpha(actual, rhs) def test_reduce_equality_mask_image_domain_symbol_side_noop(): @@ -849,49 +846,47 @@ def test_reduce_equality_mask_image_domain_symbol_side_noop(): X, Y = backend.define_vars("X", "Y", ret="stream") g = backend.define_vars("g", arg_types=(backend.scalar_typ,), ret="scalar") - term = Min.reduce(Min.mask(g(x()), x() == y()), {x: X(), y: Y()}) - with handler(ReduceEqualityMaskRange()): - actual = evaluate(term) - assert syntactic_eq_alpha(actual, term) + lhs = Min.reduce(Min.mask(g(x()), x() == y()), {x: X(), y: Y()}) + backend.check_rewrite( + lhs=lhs, rhs=lhs, rule=coproduct(ReduceEqualityMaskRange(), PlusCastIterable()) + ) @pytest.mark.parametrize("monoid", ALL_MONOIDS) def test_reduce_equality_mask_noncanonical_range(backend: Backend, monoid): """Equality elimination itself is independent of the domain representation.""" - a, c = backend.define_vars("a", "c", ret="scalar") - f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="scalar") + a, c = backend.define_vars("a", "c", ret=int) - lhs = monoid.reduce(monoid.mask(f(a()), a() == c()), {a: range(1, 4)}) + lhs = monoid.reduce(monoid.mask(a(), a() == c()), {a: range(1, 4)}) rhs = monoid.reduce( - monoid.mask(f(a()), And.plus()), - {a: Intersection.plus(as_relation(range(1, 4)), as_relation([c()]))}, + monoid.mask(a(), And.plus()), + {a: Intersection.plus(as_iterable(range(1, 4)), as_iterable([c()]))}, + ) + backend.check_rewrite( + lhs=lhs, rhs=rhs, rule=coproduct(ReduceEqualityMaskRange(), PlusCastIterable()) ) - with handler(ReduceEqualityMaskRange()): - actual = evaluate(lhs) - assert syntactic_eq_alpha(actual, rhs) @pytest.mark.parametrize("monoid", ALL_MONOIDS) def test_reduce_intersection_singleton_range(backend: Backend, monoid): """A simple-range intersection lowers to a bounds-guarded singleton.""" - x, y = backend.define_vars("x", "y", ret="scalar") - f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="scalar") + x, y = backend.define_vars("x", "y", ret=int) lhs = monoid.reduce( - f(x()), - {x: Intersection.plus(as_relation(range(3)), as_relation([y()]))}, + x(), {x: Intersection.plus(as_iterable(range(3)), as_iterable([y()]))} ) - rhs = monoid.reduce( - monoid.mask(f(x()), And.plus(0 <= y(), y() < 3)), - {x: (y(),)}, + rhs = monoid.reduce(monoid.mask(x(), And.plus(0 <= y(), y() < 3)), {x: (y(),)}) + backend.check_rewrite( + lhs=lhs, + rhs=rhs, + rule=coproduct(ReduceIntersectionSingletonRange(), PlusCastIterable()), ) - backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceIntersectionSingletonRange()) @pytest.mark.parametrize("monoid", ALL_MONOIDS) def test_reduce_equality_mask_plus(backend: Backend, monoid): """Targeted splitting exposes an intersection in the matching summand.""" - a, c = backend.define_vars("a", "c", ret="scalar") + a, c = backend.define_vars("a", "c", ret=int) f, g = backend.define_vars("f", "g", arg_types=(backend.scalar_typ,), ret="scalar") body = monoid.plus( @@ -902,11 +897,13 @@ def test_reduce_equality_mask_plus(backend: Backend, monoid): rhs = monoid.plus( monoid.reduce( monoid.mask(f(a()), And.plus()), - {a: Intersection.plus(as_relation(range(3)), as_relation([c()]))}, + {a: Intersection.plus(as_iterable(range(3)), as_iterable([c()]))}, ), monoid.reduce(monoid.mask(g(a()), c() == 0), {a: range(3)}), ) - backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceEqualityMaskRange()) + backend.check_rewrite( + lhs=lhs, rhs=rhs, rule=coproduct(ReduceEqualityMaskRange(), PlusCastIterable()) + ) def test_reduce_independent_1(backend: Backend): From 8a2c9b92a4802e4039ba34f5c1ed67dde23d9a0c Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Mon, 3 Aug 2026 15:50:24 -0400 Subject: [PATCH 04/15] lint --- tests/_monoid_helpers.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/_monoid_helpers.py b/tests/_monoid_helpers.py index 91f74cb13..ba194cea2 100644 --- a/tests/_monoid_helpers.py +++ b/tests/_monoid_helpers.py @@ -168,6 +168,14 @@ def __init__(self): def eq(self, a: Any, b: Any) -> bool: raise NotImplementedError + @staticmethod + @abstractmethod + def strategy( + arg_types: tuple[type, ...] = (), + ret: Literal["scalar", "stream"] | type = "scalar", + ) -> SearchStrategy: + raise NotImplementedError + def _fresh_op( self, name: str, From fc5ee30afee0db9508d2acf4852ca1723df27fb1 Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Tue, 11 Aug 2026 12:12:29 -0400 Subject: [PATCH 05/15] wip --- effectful/ops/monoid.py | 46 ++++++++++++++--------------------------- 1 file changed, 15 insertions(+), 31 deletions(-) diff --git a/effectful/ops/monoid.py b/effectful/ops/monoid.py index 1298d4bfd..9abd030c3 100644 --- a/effectful/ops/monoid.py +++ b/effectful/ops/monoid.py @@ -250,7 +250,9 @@ def __call__(self, t: T) -> bool: return t in self.elems -is_commutative = _ExtensiblePredicate({Max, Min, Sum, Product, And, Or}) +is_commutative = _ExtensiblePredicate( + {Max, Min, Sum, Product, And, Or, Union, Intersection} +) is_idempotent = _ExtensiblePredicate({Max, Min, And, Or, Intersection}) @@ -1442,15 +1444,21 @@ class CartesianProductPlus(ObjectInterpretation): @implements(CartesianProduct.plus) def plus(self, *args): - if not args: - return fwd() - if any(isinstance(x, Term) for x in args): - return fwd() - if not all(isinstance(x, Iterable) for x in args): + args = tuple(_unwrap_as_iterable(arg) for arg in args) + if not args or any(isinstance(x, Term) for x in args): return fwd() return [_disjoint_merge(*vals) for vals in itertools.product(*args)] +class UnionPlus(ObjectInterpretation): + @implements(Union.plus) + def plus(self, *args): + args = tuple(_unwrap_as_iterable(arg) for arg in args) + if not args or any(isinstance(x, Term) for x in args): + return fwd() + return list(itertools.chain(*args)) + + class PlusCastIterable(ObjectInterpretation): """Cast heterogeneous iterable arguments to a common iterable type.""" @@ -1465,30 +1473,6 @@ def _plus(self, monoid, *args): return monoid.plus(*(as_iterable(arg) for arg in args)) -class IterablePlus(ObjectInterpretation): - """Pure-Python filtering implementation of :data:`Intersection`. - - This preserves occurrences from the leftmost stream. Array-valued elements - remain symbolic so backend-specific, pointwise intersection lowering can - handle them. - """ - - @implements(Intersection.plus) - def _intersection_plus(self, *args): - args = tuple(_unwrap_as_iterable(arg) for arg in args) - if not args or any(isinstance(x, Term) for x in args): - return fwd() - - values = list(args[0]) - tails = [list(arg) for arg in args[1:]] - if not all( - self._concrete_value(value) for value in itertools.chain(values, *tails) - ): - return fwd() - - return list(itertools.chain(*args)) - - is_scalar = _ExtensiblePredicate({Min, Max, Sum, Product, And, Or}) @@ -1957,7 +1941,7 @@ def extend(self, *intps: Interpretation) -> typing.Self: ArgMinPlus(), ArgMaxPlus(), CartesianProductPlus(), - IterablePlus(), + UnionPlus(), ReduceEqualityMaskRange(), ReduceIntersectionSingletonRange(), ReduceWhereToMasks(), From c2fd84a1b3ad9851ddac2bdf0b1af2dce4044ec1 Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Tue, 11 Aug 2026 12:18:29 -0400 Subject: [PATCH 06/15] wip --- effectful/ops/monoid.py | 240 ++++++++++++++++++++++++++++------------ 1 file changed, 167 insertions(+), 73 deletions(-) diff --git a/effectful/ops/monoid.py b/effectful/ops/monoid.py index 9abd030c3..4e7693a27 100644 --- a/effectful/ops/monoid.py +++ b/effectful/ops/monoid.py @@ -367,12 +367,21 @@ def _(self, monoid, value, cond): is_equality = _ExtensiblePredicate({_NumberTerm.__eq__}) -def solve_group_equality(equality: Term[bool], index: int) -> Term[bool]: +def solve_group_equality( + equality: Term[bool], + index: int, + *, + side: typing.Literal["left", "right"] | None = None, +) -> Term[bool]: """Isolate one argument of a group ``plus`` in an equality. Given ``G.plus(a0, ..., an) == b``, isolate the argument at ``index``. - Negative indices follow normal Python indexing. The group expression may - occur on either side of the equality. + Negative indices follow normal Python indexing. ``side`` selects which + side to isolate when both sides are group expressions; when omitted, the + leftmost side that is a group expression is selected. + + The result preserves operand order and is therefore valid for + noncommutative groups. """ if not ( isinstance(equality, Term) @@ -383,6 +392,8 @@ def solve_group_equality(equality: Term[bool], index: int) -> Term[bool]: raise TypeError("expected an equality Term") if not isinstance(index, int): raise TypeError("argument index must be an integer") + if side not in (None, "left", "right"): + raise ValueError("side must be 'left' or 'right'") left, right = equality.args @@ -392,13 +403,23 @@ def group_plus(value): monoid = value.op.__self__ return monoid if isinstance(monoid, Group) else None - monoid = group_plus(left) - if monoid is not None: - plus_term, other = left, right - elif (monoid := group_plus(right)) is not None: - plus_term, other = right, left + left_group = group_plus(left) + right_group = group_plus(right) + if ( + left_group is not None + and right_group is not None + and left_group is not right_group + ): + raise TypeError("equality sides must use the same Group.plus") + + if side == "left" or (side is None and left_group is not None): + monoid, plus_term, other = left_group, left, right + elif side == "right" or (side is None and right_group is not None): + monoid, plus_term, other = right_group, right, left else: raise TypeError("expected an equality containing a Group.plus Term") + if monoid is None: + raise TypeError(f"expected {side} side to be a Group.plus Term") assert isinstance(plus_term, Term) try: @@ -416,8 +437,84 @@ def group_plus(value): return equality.op(target, isolated) +@dataclass(frozen=True) +class _EqualitySolution: + stream_op: Operation + rhs: Expr + + +def _solve_stream_equality( + equality: Expr[bool], streams: Streams +) -> tuple[_EqualitySolution, ...]: + """Return all group-law solutions for bare stream operands in an equality.""" + if not ( + isinstance(equality, Term) + and is_equality(equality.op) + and len(equality.args) == 2 + and not equality.kwargs + ): + return () + + def stream_symbol(value): + if ( + isinstance(value, Term) + and not value.args + and not value.kwargs + and value.op in streams + ): + return value.op + return None + + left, right = equality.args + solutions: list[_EqualitySolution] = [] + + # Direct equalities do not need a group. Preserve the existing behavior of + # leaving equalities between two bare stream symbols for a separate rule. + for stream_term, expr in ((left, right), (right, left)): + stream_op = stream_symbol(stream_term) + if ( + stream_op is not None + and stream_symbol(expr) is None + and stream_op not in fvsof(expr) + ): + solutions.append(_EqualitySolution(stream_op, expr)) + + def group_of(value): + if not (isinstance(value, Term) and _is_monoid_plus(value.op)): + return None + monoid = value.op.__self__ + return monoid if isinstance(monoid, Group) else None + + left_group, right_group = group_of(left), group_of(right) + if ( + left_group is not None + and right_group is not None + and left_group is not right_group + ): + return tuple(solutions) + + sides: tuple[tuple[typing.Literal["left", "right"], Expr], ...] = ( + ("left", left), + ("right", right), + ) + for side, value in sides: + if group_of(value) is None: + continue + assert isinstance(value, Term) + for index, operand in enumerate(value.args): + stream_op = stream_symbol(operand) + if stream_op is None: + continue + solved = solve_group_equality(equality, index, side=side) + rhs = solved.args[1] + if stream_op not in fvsof(rhs): + solutions.append(_EqualitySolution(stream_op, rhs)) + + return tuple(solutions) + + class ReduceEqualityMaskRange(ObjectInterpretation): - """Eliminate an equality by intersecting a stream with a singleton. + """Eliminate a batch of equalities by intersecting streams with singletons. For either orientation of an equality between a reduced variable and an expression that is not another reduced-variable symbol:: @@ -426,11 +523,14 @@ class ReduceEqualityMaskRange(ObjectInterpretation): == M.reduce(M.mask(v, And.plus(*m)), {i: Intersection.plus(I, [x])} | S) - The expression ``x`` may depend on streams in ``S``. Keeping those streams - in the outer bundle makes the intersection pointwise, so this rewrite does - not require idempotence or source-variable liveness checks. A separate - :class:`ReduceIntersectionSingletonRange` rule lowers intersections with - simple ranges to singleton gathers guarded by bounds masks. + The expression ``x`` may also be isolated from an equality between two + ``Group.plus`` expressions. Group inverses are applied in order, so this is + valid for noncommutative groups. ``x`` may depend on streams in ``S``. + Keeping those streams in the outer bundle makes the intersection pointwise, + so this rewrite does not require idempotence or source-variable liveness + checks. A separate :class:`ReduceIntersectionSingletonRange` rule lowers + intersections with simple ranges to singleton gathers guarded by bounds + masks. When the reduce body is a ``plus`` of the same monoid, the rule distributes the reduce over the plus -- but only when doing so exposes an eliminable @@ -441,76 +541,70 @@ class ReduceEqualityMaskRange(ObjectInterpretation): """ @staticmethod - def _match_eq(cond, streams): - """Match ``stream_op == expr`` in either orientation. - - Equalities between two bare reduced-variable symbols are left alone; - choosing which domain should represent their intersection requires a - separate rule. - """ - - def is_stream_symbol(term): - return ( - isinstance(term, Term) - and not term.args - and not term.kwargs - and term.op in streams - ) + def _rhs_cost(expr) -> tuple[float, ...]: + return (len(fvsof(expr)), sizeof(expr)) - if not (isinstance(cond, Term) and is_equality(cond.op)): - return None - - lhs, rhs = cond.args - for stream_term, expr in ((lhs, rhs), (rhs, lhs)): - if is_stream_symbol(stream_term) and not is_stream_symbol(expr): - return stream_term.op, expr - return None - - def _eliminate(self, monoid, value, mask, streams): - """Discharge one eliminable equality constraint, or return ``None``.""" + @classmethod + def _eliminate(cls, monoid, value, mask, streams): conds = _conjuncts(mask) - for i, cond in enumerate(conds): - matched = self._match_eq(cond, streams) - if matched is None: + matched_conds = [ + (i, solution.stream_op, solution.rhs) + for i, cond in enumerate(conds) + for solution in _solve_stream_equality(cond, streams) + ] + ranked_conds = sorted(matched_conds, key=lambda c: (*cls._rhs_cost(c[2]), c[0])) + + # Build a compatible batch greedily. A conjunct and lhs may each occur + # only once in the batch, and none of the batch's lhs variables may + # occur in any of its rhs expressions. Other equalities remain in the + # mask for a later elimination pass. + selected = [] + selected_indices = set() + selected_lhs = set() + selected_rhs_fvs = set() + for candidate in ranked_conds: + index, stream_op, expr = candidate + expr_fvs = fvsof(expr) + if ( + index in selected_indices + or stream_op in selected_lhs + or stream_op in selected_rhs_fvs + or selected_lhs & expr_fvs + ): continue - stream_op, expr = matched + selected.append(candidate) + selected_indices.add(index) + selected_lhs.add(stream_op) + selected_rhs_fvs.update(expr_fvs) - new_mask = monoid.mask( - value, - And.plus(*(c for (j, c) in enumerate(conds) if i != j)), - ) - new_streams = { - stream_op: Intersection.plus(streams[stream_op], [expr]), - } | {k: v for (k, v) in streams.items() if k is not stream_op} - return monoid.reduce(new_mask, new_streams) - return None + if not selected: + return None - def _summand_eliminable(self, monoid, summand, streams): - return ( - isinstance(summand, Term) - and _is_monoid_mask(summand.op) - and summand.op.__self__ == monoid - and self._eliminate(monoid, summand.args[0], summand.args[1], streams) - is not None + new_mask = monoid.mask( + value, + And.plus(*(c for i, c in enumerate(conds) if i not in selected_indices)), ) + replacements = { + stream_op: Intersection.plus(streams[stream_op], [expr]) + for _, stream_op, expr in selected + } + new_streams = { + stream_op: replacements.get(stream_op, stream) + for stream_op, stream in streams.items() + } + return monoid.reduce(new_mask, new_streams) @implements(Monoid.reduce) def _(self, monoid, body, streams): - if not isinstance(body, Term): + if not ( + isinstance(body, Term) + and _is_monoid_mask(body.op) + and body.op.__self__ == monoid + ): return fwd() - # single mask body: discharge an equality constraint directly - if _is_monoid_mask(body.op) and body.op.__self__ == monoid: - result = self._eliminate(monoid, body.args[0], body.args[1], streams) - return result if result is not None else fwd() - - # plus body: distribute the reduce only when it exposes an eliminable - # equality mask in some summand - if _is_monoid_plus(body.op) and body.op.__self__ == monoid: - if any(self._summand_eliminable(monoid, s, streams) for s in body.args): - return monoid.plus(*(monoid.reduce(s, streams) for s in body.args)) - - return fwd() + result = self._eliminate(monoid, body.args[0], body.args[1], streams) + return result if result is not None else fwd() class ReduceIntersectionSingletonRange(ObjectInterpretation): From cad1c41c9e0c9fa0694fa2d690f1b18c05c33902 Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Tue, 11 Aug 2026 12:19:27 -0400 Subject: [PATCH 07/15] wip --- effectful/ops/monoid.py | 10 +++++++++- effectful/ops/semantics.py | 34 ++++++++++++++++++++++++++++++++++ 2 files changed, 43 insertions(+), 1 deletion(-) diff --git a/effectful/ops/monoid.py b/effectful/ops/monoid.py index 4e7693a27..bede4a8dc 100644 --- a/effectful/ops/monoid.py +++ b/effectful/ops/monoid.py @@ -19,7 +19,15 @@ import effectful.ops.syntax from effectful.internals.runtime import interpreter -from effectful.ops.semantics import coproduct, evaluate, fvsof, fwd, handler, typeof +from effectful.ops.semantics import ( + coproduct, + evaluate, + fvsof, + fwd, + handler, + sizeof, + typeof, +) from effectful.ops.syntax import ( ObjectInterpretation, Scoped, diff --git a/effectful/ops/semantics.py b/effectful/ops/semantics.py index be3686dd6..1af1ec580 100644 --- a/effectful/ops/semantics.py +++ b/effectful/ops/semantics.py @@ -427,3 +427,37 @@ def fvsof[S](term: Expr[S]) -> collections.abc.Set[Operation]: """ result = evaluate(term, intp=_FVSOF_INTP) return result.fvs if isinstance(result, _FvsAnalysis) else frozenset() + + +class _SizeofAnalysisValue: + value: int + + def __init__(self, value: int | typing.Self): + self.value = value.value if isinstance(value, _SizeofAnalysisValue) else value + + def __add__(self, other): + if isinstance(other, _SizeofAnalysisValue): + return _SizeofAnalysisValue(self.value + other.value) + return NotImplemented + + __radd__ = __add__ + + +@functools.cache +def _sizeof_intp(): + def _apply(_, *args, **kwargs): + s_args = [ + x if isinstance(x, _SizeofAnalysisValue) else _SizeofAnalysisValue(1) + for x in (*args, *kwargs.values()) + ] + size = sum(s_args, start=_SizeofAnalysisValue(1)) + return size + + return {apply: _apply} + + +def sizeof(term: Expr) -> int: + result = evaluate(term, intp=_sizeof_intp()) + if not isinstance(result, _SizeofAnalysisValue): + return 1 + return result.value From abf26865958aa2b543f15a78a46dc0cd16d85ea9 Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Tue, 11 Aug 2026 14:38:01 -0400 Subject: [PATCH 08/15] fix bugs --- effectful/ops/monoid.py | 55 +++++++++++++++++++++++++++++++--------- tests/test_ops_monoid.py | 45 ++++++++++++++++++++++++++++++++ 2 files changed, 88 insertions(+), 12 deletions(-) diff --git a/effectful/ops/monoid.py b/effectful/ops/monoid.py index bede4a8dc..e25026eea 100644 --- a/effectful/ops/monoid.py +++ b/effectful/ops/monoid.py @@ -258,10 +258,8 @@ def __call__(self, t: T) -> bool: return t in self.elems -is_commutative = _ExtensiblePredicate( - {Max, Min, Sum, Product, And, Or, Union, Intersection} -) -is_idempotent = _ExtensiblePredicate({Max, Min, And, Or, Intersection}) +is_commutative = _ExtensiblePredicate({Max, Min, Sum, Product, And, Or, Union}) +is_idempotent = _ExtensiblePredicate({Max, Min, And, Or}) @dataclass @@ -539,19 +537,27 @@ class ReduceEqualityMaskRange(ObjectInterpretation): checks. A separate :class:`ReduceIntersectionSingletonRange` rule lowers intersections with simple ranges to singleton gathers guarded by bounds masks. - - When the reduce body is a ``plus`` of the same monoid, the rule distributes - the reduce over the plus -- but only when doing so exposes an eliminable - equality mask in some summand. This is a *targeted* split (it leaves - ``ReduceSplit`` conservative): summands whose reduced index appears in an - equality get restricted domains, while the rest stay as ordinary masked - reduces. """ @staticmethod def _rhs_cost(expr) -> tuple[float, ...]: return (len(fvsof(expr)), sizeof(expr)) + @staticmethod + def _can_restrict(stream) -> bool: + """Whether an equality may add the first restriction to ``stream``. + + A stream already represented by an intersection, or by the singleton + produced when lowering one, has already had a substitution variable + selected. Leaving further equalities in the mask allows singleton + substitution to turn them into residual conditions instead of building + nested intersections. + """ + stream = _unwrap_as_iterable(stream) + if isinstance(stream, Term) and stream.op is Intersection.plus: + return False + return not (isinstance(stream, Sequence) and len(stream) == 1) + @classmethod def _eliminate(cls, monoid, value, mask, streams): conds = _conjuncts(mask) @@ -559,13 +565,15 @@ def _eliminate(cls, monoid, value, mask, streams): (i, solution.stream_op, solution.rhs) for i, cond in enumerate(conds) for solution in _solve_stream_equality(cond, streams) + if cls._can_restrict(streams[solution.stream_op]) ] ranked_conds = sorted(matched_conds, key=lambda c: (*cls._rhs_cost(c[2]), c[0])) # Build a compatible batch greedily. A conjunct and lhs may each occur # only once in the batch, and none of the batch's lhs variables may # occur in any of its rhs expressions. Other equalities remain in the - # mask for a later elimination pass. + # mask; once a selected stream becomes an intersection or singleton, + # they are retained until substitution turns them into residual masks. selected = [] selected_indices = set() selected_lhs = set() @@ -1561,6 +1569,28 @@ def plus(self, *args): return list(itertools.chain(*args)) +class IntersectionPlus(ObjectInterpretation): + """Pure-Python implementation of :data:`Intersection`. + + The result preserves the order and multiplicity of the first iterable and + retains each of its elements that occurs in every subsequent iterable. + """ + + @implements(Intersection.plus) + def plus(self, *args): + args = tuple(_unwrap_as_iterable(arg) for arg in args) + if ( + not args + or any(isinstance(arg, Term) for arg in args) + or any(fvsof(arg) for arg in args) + ): + return fwd() + + first, *rest = args + rest = [tuple(values) for values in rest] + return [value for value in first if all(value in values for values in rest)] + + class PlusCastIterable(ObjectInterpretation): """Cast heterogeneous iterable arguments to a common iterable type.""" @@ -2044,6 +2074,7 @@ def extend(self, *intps: Interpretation) -> typing.Self: ArgMaxPlus(), CartesianProductPlus(), UnionPlus(), + IntersectionPlus(), ReduceEqualityMaskRange(), ReduceIntersectionSingletonRange(), ReduceWhereToMasks(), diff --git a/tests/test_ops_monoid.py b/tests/test_ops_monoid.py index b83cb11ce..fefb758b7 100644 --- a/tests/test_ops_monoid.py +++ b/tests/test_ops_monoid.py @@ -19,6 +19,7 @@ Factor, Group, Intersection, + IntersectionPlus, InverseInverse, InversePlus, Max, @@ -58,6 +59,7 @@ as_iterable, distributes_over, is_commutative, + is_idempotent, solve_group_equality, ) from effectful.ops.semantics import coproduct, evaluate, fvsof, handler @@ -867,6 +869,20 @@ def test_reduce_equality_mask_noncanonical_range(backend: Backend, monoid): ) +def test_reduce_equality_mask_repeated_constraints_eliminate_intersections(): + """Repeated equalities for one stream still normalize to supported domains.""" + x, a, b = IntBackend().define_vars("x", "a", "b", ret=int) + lhs = Sum.reduce( + Sum.mask(x(), And.plus(x() == a(), x() == b())), + {x: range(3)}, + ) + + with handler(NormalizeIntp), handler(EvaluateIntp): + result = evaluate(lhs) + + assert Intersection.plus not in fvsof(result) + + @pytest.mark.parametrize("monoid", ALL_MONOIDS) def test_reduce_intersection_singleton_range(backend: Backend, monoid): """A simple-range intersection lowers to a bounds-guarded singleton.""" @@ -1332,6 +1348,35 @@ def _f(v: int) -> float: assert math.isclose(result, 10.0) +def test_intersection_plus_preserves_left_order_and_multiplicity(): + with handler(IntersectionPlus()): + result = Intersection.plus( + as_iterable([3, 1, 1, 2, 4]), + as_iterable([1, 2, 3]), + as_iterable([1, 3]), + ) + + assert result == [3, 1, 1] + + +def test_intersection_is_not_registered_commutative_or_idempotent(): + assert not is_commutative(Intersection) + assert not is_idempotent(Intersection) + + with handler(IntersectionPlus()): + assert Intersection.plus([1, 1], [1]) == [1, 1] + assert Intersection.plus([1], [1, 1]) == [1] + + +def test_equality_intersection_preserves_stream_multiplicity(): + x = Operation.define(int, name="x") + + with handler(NormalizeIntp), handler(EvaluateIntp): + result = evaluate(Sum.reduce(Sum.mask(x(), x() == 1), {x: [1, 1, 2]})) + + assert result == 2 + + # --------------------------------------------------------------------------- # CartesianProduct.plus (pure-Python ``CartesianProductPlus`` implementation) # --------------------------------------------------------------------------- From 0001ed6b487c8e307fb76b52baa5a87dd3f43220 Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Tue, 11 Aug 2026 16:14:00 -0400 Subject: [PATCH 09/15] fix tests --- effectful/ops/monoid.py | 2 +- tests/test_ops_monoid.py | 23 ----------------------- 2 files changed, 1 insertion(+), 24 deletions(-) diff --git a/effectful/ops/monoid.py b/effectful/ops/monoid.py index e25026eea..bbe0df5b9 100644 --- a/effectful/ops/monoid.py +++ b/effectful/ops/monoid.py @@ -224,7 +224,7 @@ def __init__(self, name: str, identity: T, zero: T): @Operation.define -def as_iterable[T](value: Iterable[T]) -> Iterable[T]: +def as_iterable(value: Iterable) -> Iterable: raise NotHandled diff --git a/tests/test_ops_monoid.py b/tests/test_ops_monoid.py index fefb758b7..ffeb06622 100644 --- a/tests/test_ops_monoid.py +++ b/tests/test_ops_monoid.py @@ -899,29 +899,6 @@ def test_reduce_intersection_singleton_range(backend: Backend, monoid): ) -@pytest.mark.parametrize("monoid", ALL_MONOIDS) -def test_reduce_equality_mask_plus(backend: Backend, monoid): - """Targeted splitting exposes an intersection in the matching summand.""" - a, c = backend.define_vars("a", "c", ret=int) - f, g = backend.define_vars("f", "g", arg_types=(backend.scalar_typ,), ret="scalar") - - body = monoid.plus( - monoid.mask(f(a()), a() == c()), - monoid.mask(g(a()), c() == 0), - ) - lhs = monoid.reduce(body, {a: range(3)}) - rhs = monoid.plus( - monoid.reduce( - monoid.mask(f(a()), And.plus()), - {a: Intersection.plus(as_iterable(range(3)), as_iterable([c()]))}, - ), - monoid.reduce(monoid.mask(g(a()), c() == 0), {a: range(3)}), - ) - backend.check_rewrite( - lhs=lhs, rhs=rhs, rule=coproduct(ReduceEqualityMaskRange(), PlusCastIterable()) - ) - - def test_reduce_independent_1(backend: Backend): a, b = backend.define_vars("a", "b", ret="scalar") A, B = backend.define_vars("A", "B", ret="stream") From 74dc488ac4637b87194dc3dc7191e83f4f8cf7c7 Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Tue, 11 Aug 2026 16:40:44 -0400 Subject: [PATCH 10/15] push masks when factoring --- effectful/ops/monoid.py | 52 ++++++++++++++++++++++++----------------- 1 file changed, 31 insertions(+), 21 deletions(-) diff --git a/effectful/ops/monoid.py b/effectful/ops/monoid.py index bbe0df5b9..6c8dd3246 100644 --- a/effectful/ops/monoid.py +++ b/effectful/ops/monoid.py @@ -15,7 +15,7 @@ ) from dataclasses import dataclass from graphlib import TopologicalSorter -from typing import Annotated, Any +from typing import Annotated, Any, Literal import effectful.ops.syntax from effectful.internals.runtime import interpreter @@ -99,7 +99,7 @@ def inner_stream( ) -def inner_streams_first(streams: dict[Operation, Expr]) -> Iterable[Operation]: +def inner_streams_first(streams: Streams) -> Iterable[Operation]: """Iterable over streams where dependent streams precede their dependencies.""" stream_vars = set(streams.keys()) @@ -917,7 +917,7 @@ def reduce(self, monoid, body, streams): @Operation.define def choose_contraction(factors: Sequence[Any], streams: Streams) -> Operation: - """Used by `ReduceFactorization` to choose a contraction when there is + """Used by `Factor` to choose a contraction when there is ambiguity. Takes the factors and streams that are eligible for contraction (innermost and non-universal). @@ -940,8 +940,20 @@ def choose_contraction(factors: Sequence[Any], streams: Streams) -> Operation: class Factor(ObjectInterpretation): + def mask_plus( + self, + outer_monoid: Monoid, + inner_monoid: Monoid, + *factors: tuple[Literal["factor", "mask"], Expr], + ) -> Expr: + """Turn a flat list of factors and masks into a masked plus.""" + return outer_monoid.mask( + inner_monoid.plus(*(f for (k, f) in factors if k == "factor")), + And.plus(*(f for (k, f) in factors if k == "mask")), + ) + @implements(Monoid.reduce) - def reduce(self, monoid, body, streams): + def reduce(self, monoid: Monoid, body, streams: Streams): """reduce(⊗(F_v ∪ F_rest), {v} ∪ S) = reduce(⊗F_rest ⊗ reduce(⊗F_v, {v}), S) where F_v = factors mentioning v, F_rest = the others. Fires only when @@ -969,6 +981,7 @@ def reduce(self, monoid, body, streams): plus_term = body if _is_monoid_mask(body.op) and body.op.__self__ is monoid: plus_term, cond = body.args + conds = _conjuncts(cond) if not ( isinstance(plus_term, Term) @@ -979,8 +992,9 @@ def reduce(self, monoid, body, streams): inner = plus_term.op.__self__ stream_keys = set(streams) - cond_fvs = fvsof(cond) if cond is not None else set() - factors = [(a, fvsof(a) | cond_fvs) for a in plus_term.args] + factors: list[tuple[Literal["factor", "mask"], Expr]] = [ + ("factor", a) for a in plus_term.args + ] + [("mask", c) for c in conds] # candidates: innermost-eligible (no remaining stream depends on v), # non-universal (some factor doesn't mention v) @@ -988,9 +1002,7 @@ def reduce(self, monoid, body, streams): for k, v in streams.items(): if any(k in fvsof(vv) for kk, vv in streams.items() if k is not kk): continue - if len({i for i, (_, fvs) in enumerate(factors) if k in fvs}) == len( - factors - ): + if all(k in fvsof(factor) for (_, factor) in factors): continue # v is universal: leave it in the outer core eligible[k] = v @@ -1002,22 +1014,20 @@ def reduce(self, monoid, body, streams): inner_stream = choose_contraction(plus_term.args, eligible) inner_factor_ids = frozenset( - i for i, (_, fvs) in enumerate(factors) if inner_stream in fvs + i for i, (_, factor) in enumerate(factors) if inner_stream in fvsof(factor) ) - inner_factors = [factors[i][0] for i in sorted(inner_factor_ids)] + inner_factors = [factors[i] for i in sorted(inner_factor_ids)] inner_stream_keys = {inner_stream} inner_deps = set().union( - *(factors[i][1] for i in inner_factor_ids), + *(fvsof(factors[i][1]) for i in inner_factor_ids), fvsof(streams[inner_stream]) & stream_keys, ) - outer_factors = [ - a for i, (a, _) in enumerate(factors) if i not in inner_factor_ids - ] + outer_factors = [a for i, a in enumerate(factors) if i not in inner_factor_ids] outer_stream_keys = stream_keys - inner_stream_keys outer_factor_deps = set().union( - *(vars for i, (_, vars) in enumerate(factors) if i not in inner_factor_ids) + *(fvsof(f) for i, (_, f) in enumerate(factors) if i not in inner_factor_ids) ) # find all streams that are used in the inner factors/streams and are @@ -1038,13 +1048,13 @@ def reduce(self, monoid, body, streams): outer_stream_keys -= {s} inner_streams = {k: v for (k, v) in streams.items() if k in inner_stream_keys} - inner_red = monoid.reduce(inner.plus(*inner_factors), inner_streams) + inner_red = monoid.reduce( + self.mask_plus(monoid, inner, *inner_factors), inner_streams + ) rest_streams = {k: s for k, s in streams.items() if k in outer_stream_keys} - new_body = inner.plus(*outer_factors, inner_red) - if cond is not None: - new_body = monoid.mask(new_body, cond) - return monoid.reduce(new_body, rest_streams) if rest_streams else new_body + new_body = self.mask_plus(monoid, inner, *outer_factors, ("factor", inner_red)) + return monoid.reduce(new_body, rest_streams) class ReduceUnfactor(ObjectInterpretation): From 7e4e592a5e0d0b79bcc30d6b248e3d89094e718a Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Tue, 11 Aug 2026 16:54:04 -0400 Subject: [PATCH 11/15] fix bugs --- effectful/ops/monoid.py | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/effectful/ops/monoid.py b/effectful/ops/monoid.py index 6c8dd3246..af71da68e 100644 --- a/effectful/ops/monoid.py +++ b/effectful/ops/monoid.py @@ -944,13 +944,16 @@ def mask_plus( self, outer_monoid: Monoid, inner_monoid: Monoid, - *factors: tuple[Literal["factor", "mask"], Expr], + *factor_conds: tuple[Literal["factor", "mask"], Expr], ) -> Expr: """Turn a flat list of factors and masks into a masked plus.""" - return outer_monoid.mask( - inner_monoid.plus(*(f for (k, f) in factors if k == "factor")), - And.plus(*(f for (k, f) in factors if k == "mask")), - ) + factors, conds = [], [] + for k, f in factor_conds: + (factors if k == "factor" else conds).append(f) + + term = inner_monoid.plus(*factors) + term = term if not conds else outer_monoid.mask(term, And.plus(*conds)) + return term @implements(Monoid.reduce) def reduce(self, monoid: Monoid, body, streams: Streams): @@ -977,11 +980,11 @@ def reduce(self, monoid: Monoid, body, streams: Streams): return fwd() # Optionally peel an outer mask of the reduce monoid. - cond = None plus_term = body + conds = () if _is_monoid_mask(body.op) and body.op.__self__ is monoid: plus_term, cond = body.args - conds = _conjuncts(cond) + conds = _conjuncts(cond) if not ( isinstance(plus_term, Term) @@ -1054,7 +1057,7 @@ def reduce(self, monoid: Monoid, body, streams: Streams): rest_streams = {k: s for k, s in streams.items() if k in outer_stream_keys} new_body = self.mask_plus(monoid, inner, *outer_factors, ("factor", inner_red)) - return monoid.reduce(new_body, rest_streams) + return monoid.reduce(new_body, rest_streams) if rest_streams else new_body class ReduceUnfactor(ObjectInterpretation): From 935ecc9a7b261ca9da9af04fbccc5a120e0bad65 Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Tue, 11 Aug 2026 17:23:47 -0400 Subject: [PATCH 12/15] add tests --- tests/test_ops_monoid.py | 32 ++++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/tests/test_ops_monoid.py b/tests/test_ops_monoid.py index ffeb06622..3f2f08de7 100644 --- a/tests/test_ops_monoid.py +++ b/tests/test_ops_monoid.py @@ -1008,6 +1008,38 @@ def test_reduce_lift_shared(outer, inner, backend: Backend): backend.check_rewrite(lhs=lhs, rhs=rhs, rule=Factor()) +def test_reduce_factor_mask(backend: Backend): + a, b, c, d = backend.define_vars("a", "b", "c", "d", ret="scalar") + A, B = backend.define_vars("A", "B", ret="stream") + + lhs = Sum.reduce( + Sum.mask(Product.plus(a(), b()), And.plus(a() == c(), b() == d())), + {a: A(), b: B()}, + ) + rhs = Product.plus( + Sum.reduce(Sum.mask(Product.plus(a()), And.plus(a() == c())), {a: A()}), + Sum.reduce(Sum.mask(Product.plus(b()), And.plus(b() == d())), {b: B()}), + ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=Factor()) + + +def test_reduce_factor_mask_2(backend: Backend): + a, b = backend.define_vars("a", "b", ret="scalar") + A, B = backend.define_vars("A", "B", ret="stream") + + lhs = Sum.reduce( + Sum.mask(Product.plus(a(), b()), And.plus(a() == b())), + {a: A(), b: B()}, + ) + lhs = Sum.reduce( + Product.plus( + b(), Sum.reduce(Sum.mask(Product.plus(a()), And.plus(a() == b())), {a: A()}) + ), + {b: B()}, + ) + backend.check_rewrite(lhs=lhs, rhs=lhs, rule=Factor()) + + @pytest.mark.parametrize("outer,inner", MONOID_PAIRS) def test_reduce_lift_shared_deps(outer, inner, backend: Backend): """A shared stream is lifted together with its dependencies: both ``c`` From 816f50609f19d608c0285bf9769ea729bf9e45d3 Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Tue, 11 Aug 2026 17:25:34 -0400 Subject: [PATCH 13/15] rename --- effectful/ops/monoid.py | 4 ++-- tests/test_ops_monoid.py | 16 ++++++++-------- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/effectful/ops/monoid.py b/effectful/ops/monoid.py index af71da68e..e88066104 100644 --- a/effectful/ops/monoid.py +++ b/effectful/ops/monoid.py @@ -519,7 +519,7 @@ def group_of(value): return tuple(solutions) -class ReduceEqualityMaskRange(ObjectInterpretation): +class ReduceEqualityMask(ObjectInterpretation): """Eliminate a batch of equalities by intersecting streams with singletons. For either orientation of an equality between a reduced variable and an @@ -2088,7 +2088,7 @@ def extend(self, *intps: Interpretation) -> typing.Self: CartesianProductPlus(), UnionPlus(), IntersectionPlus(), - ReduceEqualityMaskRange(), + ReduceEqualityMask(), ReduceIntersectionSingletonRange(), ReduceWhereToMasks(), ) diff --git a/tests/test_ops_monoid.py b/tests/test_ops_monoid.py index 3f2f08de7..21fb8bc7c 100644 --- a/tests/test_ops_monoid.py +++ b/tests/test_ops_monoid.py @@ -42,7 +42,7 @@ ReduceDisjunctiveDisequalityMask, ReduceDistributeCartesianProduct, ReduceEmpty, - ReduceEqualityMaskRange, + ReduceEqualityMask, ReduceFusion, ReduceIntersectionSingletonRange, ReduceMaskHoist, @@ -756,7 +756,7 @@ def test_reduce_equality_mask_range_simple(backend: Backend, monoid): {a: Intersection.plus(as_iterable(range(3)), as_iterable([c()]))}, ) backend.check_rewrite( - lhs=lhs, rhs=rhs, rule=coproduct(ReduceEqualityMaskRange(), PlusCastIterable()) + lhs=lhs, rhs=rhs, rule=coproduct(ReduceEqualityMask(), PlusCastIterable()) ) @@ -774,7 +774,7 @@ def test_reduce_equality_mask_range_residual_conjuncts(backend: Backend, monoid) {a: Intersection.plus(as_iterable(range(4)), as_iterable([c()]))}, ) backend.check_rewrite( - lhs=lhs, rhs=rhs, rule=coproduct(ReduceEqualityMaskRange(), PlusCastIterable()) + lhs=lhs, rhs=rhs, rule=coproduct(ReduceEqualityMask(), PlusCastIterable()) ) @@ -799,7 +799,7 @@ def test_reduce_equality_mask_dependent_intersection(): ) backend.check_rewrite( - lhs=lhs, rhs=rhs, rule=coproduct(ReduceEqualityMaskRange(), PlusCastIterable()) + lhs=lhs, rhs=rhs, rule=coproduct(ReduceEqualityMask(), PlusCastIterable()) ) @@ -817,7 +817,7 @@ def test_reduce_equality_mask_dependent_intersection_nonidempotent(): {x: Intersection.plus(as_iterable(X()), as_iterable([f(y())])), y: Y()}, ) backend.check_rewrite( - lhs=lhs, rhs=rhs, rule=coproduct(ReduceEqualityMaskRange(), PlusCastIterable()) + lhs=lhs, rhs=rhs, rule=coproduct(ReduceEqualityMask(), PlusCastIterable()) ) @@ -837,7 +837,7 @@ def test_reduce_equality_mask_dependent_intersection_retains_source_uses(): {x: Intersection.plus(as_iterable(X()), as_iterable([f(y())])), y: Y()}, ) backend.check_rewrite( - lhs=lhs, rhs=rhs, rule=coproduct(ReduceEqualityMaskRange(), PlusCastIterable()) + lhs=lhs, rhs=rhs, rule=coproduct(ReduceEqualityMask(), PlusCastIterable()) ) @@ -850,7 +850,7 @@ def test_reduce_equality_mask_image_domain_symbol_side_noop(): lhs = Min.reduce(Min.mask(g(x()), x() == y()), {x: X(), y: Y()}) backend.check_rewrite( - lhs=lhs, rhs=lhs, rule=coproduct(ReduceEqualityMaskRange(), PlusCastIterable()) + lhs=lhs, rhs=lhs, rule=coproduct(ReduceEqualityMask(), PlusCastIterable()) ) @@ -865,7 +865,7 @@ def test_reduce_equality_mask_noncanonical_range(backend: Backend, monoid): {a: Intersection.plus(as_iterable(range(1, 4)), as_iterable([c()]))}, ) backend.check_rewrite( - lhs=lhs, rhs=rhs, rule=coproduct(ReduceEqualityMaskRange(), PlusCastIterable()) + lhs=lhs, rhs=rhs, rule=coproduct(ReduceEqualityMask(), PlusCastIterable()) ) From f517bdb8be254c50a9436a318062026cd23c2874 Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Tue, 11 Aug 2026 17:27:31 -0400 Subject: [PATCH 14/15] make equality mask elimination part of normalization --- effectful/ops/monoid.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/effectful/ops/monoid.py b/effectful/ops/monoid.py index e88066104..31a32da55 100644 --- a/effectful/ops/monoid.py +++ b/effectful/ops/monoid.py @@ -2088,8 +2088,6 @@ def extend(self, *intps: Interpretation) -> typing.Self: CartesianProductPlus(), UnionPlus(), IntersectionPlus(), - ReduceEqualityMask(), - ReduceIntersectionSingletonRange(), ReduceWhereToMasks(), ) @@ -2122,6 +2120,8 @@ def extend(self, *intps: Interpretation) -> typing.Self: ReduceFusion(), ReduceUnion(), ReduceSplit(), + ReduceEqualityMask(), + ReduceIntersectionSingletonRange(), Factor(), ReduceDistributeCartesianProduct(), ReduceWeightedStream(), From 7a3a0188df33f3c3a2d4613ee468cfed6b441c77 Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Tue, 11 Aug 2026 17:30:33 -0400 Subject: [PATCH 15/15] lint --- effectful/ops/monoid.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/effectful/ops/monoid.py b/effectful/ops/monoid.py index 31a32da55..f2cfee75b 100644 --- a/effectful/ops/monoid.py +++ b/effectful/ops/monoid.py @@ -234,7 +234,7 @@ def _unwrap_as_iterable[T](arg: Iterable[T]) -> Iterable[T]: return arg -def _conjuncts(mask) -> Sequence[Term]: +def _conjuncts(mask) -> Sequence[Expr]: """Return the conjuncts of an ``And`` mask as a flat tuple.""" match mask: case Term(And.plus, elems, {}): @@ -947,7 +947,8 @@ def mask_plus( *factor_conds: tuple[Literal["factor", "mask"], Expr], ) -> Expr: """Turn a flat list of factors and masks into a masked plus.""" - factors, conds = [], [] + factors: list[Expr] = [] + conds: list[Expr] = [] for k, f in factor_conds: (factors if k == "factor" else conds).append(f) @@ -981,7 +982,7 @@ def reduce(self, monoid: Monoid, body, streams: Streams): # Optionally peel an outer mask of the reduce monoid. plus_term = body - conds = () + conds: Sequence[Term] = () if _is_monoid_mask(body.op) and body.op.__self__ is monoid: plus_term, cond = body.args conds = _conjuncts(cond)