From 3b1472b2ad627c815ff4b1c11d1ffb1565b1956f Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Wed, 6 May 2026 16:21:46 -0400 Subject: [PATCH 01/16] Add monoid module (#653) * add monoid module * clean up * fix doctest * fix * wip * remove incorrect rule * add disjoint set tests and fix bug * lint * drop jax monoid defs * drop incorrect comment * add assert * reduce nondeterminism and add assertions * fix inconsistent stream numbering and missing constant factors --- effectful/internals/disjoint_set.py | 99 +++++ effectful/ops/monoid.py | 556 +++++++++++++++++++++++++++ effectful/ops/syntax.py | 78 ++++ pyproject.toml | 1 + tests/_monoid_helpers.py | 85 ++++ tests/test_internals_disjoint_set.py | 124 ++++++ tests/test_ops_monoid.py | 518 +++++++++++++++++++++++++ 7 files changed, 1461 insertions(+) create mode 100644 effectful/internals/disjoint_set.py create mode 100644 effectful/ops/monoid.py create mode 100644 tests/_monoid_helpers.py create mode 100644 tests/test_internals_disjoint_set.py create mode 100644 tests/test_ops_monoid.py diff --git a/effectful/internals/disjoint_set.py b/effectful/internals/disjoint_set.py new file mode 100644 index 000000000..73b5c5c52 --- /dev/null +++ b/effectful/internals/disjoint_set.py @@ -0,0 +1,99 @@ +class DisjointSet: + """Disjoint Set Union (Union-Find) data structure. + + Maintains a collection of disjoint sets over the integers 0..n-1, + supporting near-constant-time union and find operations via + path compression and union by rank. + + The amortized time complexity per operation is O(α(n)), where α + is the inverse Ackermann function (effectively constant for any + practical n). + + Example: + >>> dsu = DisjointSet(5) + >>> dsu.union(0, 1) + True + >>> dsu.union(1, 2) + True + >>> dsu.find(0) == dsu.find(2) + True + >>> dsu.find(0) == dsu.find(3) + False + """ + + def __init__(self, n): + """Initialize n singleton sets: {0}, {1}, ..., {n-1}. + + Args: + n: The number of elements. Elements are labeled 0..n-1. + """ + self.parent = list(range(n)) + self.rank = [0] * n + + def _validate(self, x): + if x < 0 or x >= len(self.parent): + raise IndexError(f"Element {x} out of bounds") + + def find(self, x): + """Return the representative (root) of the set containing x. + + Two elements belong to the same set if and only if they have + the same representative. Applies path compression: every node + traversed is re-parented directly to its grandparent, flattening + the tree to speed up future queries. + + Args: + x: The element to look up. + + Returns: + The root element of x's set. + """ + self._validate(x) + while self.parent[x] != x: + self.parent[x] = self.parent[self.parent[x]] # path compression + x = self.parent[x] + return x + + def union(self, *elements): + """Merge the sets containing all given elements into one. + + Accepts any number of elements and unions them all together. + Uses union by rank: shallower trees are attached under the root + of the deeper one, keeping the combined tree shallow. + + Args: + *elements: Two or more elements to merge into a single set. + Calling with 0 or 1 elements is a no-op and returns False. + + Returns: + True if any merging occurred (i.e., at least two of the + elements were in different sets); False if all elements + were already in the same set or fewer than 2 were given. + """ + if len(elements) < 2: + return False + + merged = False + first = elements[0] + + for y in elements[1:]: + if self._union_pair(first, y): + merged = True + + return merged + + def _union_pair(self, x, y): + rx = self.find(x) + ry = self.find(y) + + if rx == ry: + return False + + if self.rank[rx] < self.rank[ry]: + rx, ry = ry, rx + + self.parent[ry] = rx + if self.rank[rx] == self.rank[ry]: + self.rank[rx] += 1 + + return True diff --git a/effectful/ops/monoid.py b/effectful/ops/monoid.py new file mode 100644 index 000000000..58a10ba3d --- /dev/null +++ b/effectful/ops/monoid.py @@ -0,0 +1,556 @@ +import collections.abc +import functools +import itertools +import numbers +import typing +from collections import Counter, defaultdict +from collections.abc import Callable, Generator, Iterable, Iterator, Mapping, Sequence +from dataclasses import dataclass +from graphlib import TopologicalSorter +from typing import Annotated, Any + +from effectful.internals.disjoint_set import DisjointSet +from effectful.ops.semantics import coproduct, evaluate, fvsof, fwd, handler +from effectful.ops.syntax import ( + ObjectInterpretation, + Scoped, + _NumberTerm, + defdata, + implements, + iter_, + syntactic_eq, + syntactic_hash, +) +from effectful.ops.types import Interpretation, NotHandled, Operation, Term + +# Note: The streams value type should be something like Iterable[T], but some of +# our target stream types (e.g. jax.Array) are not subtypes of Iterable +type Streams[T] = Mapping[Operation[[], T], Any] + +type Body[T] = ( + Iterable[T] + | Callable[..., Body[T]] + | T + | Mapping[Any, Body[T]] + | Interpretation[T, Body[T]] +) + + +def order_streams[T](streams: Streams[T]) -> Iterable[tuple[Operation[[], T], Any]]: + """Determine an order to evaluate the streams based on their dependencies""" + stream_vars = set(streams.keys()) + dependencies = {k: fvsof(v) & stream_vars for k, v in streams.items()} + topo = TopologicalSorter(dependencies) + topo.prepare() + while topo.is_active(): + node_group = topo.get_ready() + for op in sorted(node_group): + yield (op, streams[op]) + topo.done(*node_group) + + +class Monoid[T]: + kernel: Operation[[T, T], T] + identity: T + + def __init__(self, kernel: Callable[[T, T], T], identity: T): + self.identity = identity + self.kernel = ( + kernel if isinstance(kernel, Operation) else Operation.define(kernel) + ) + + def __repr__(self): + return f"{type(self)}({self.kernel}, {self.identity})" + + @Operation.define + def plus[S: Body[T]](self, *args: S) -> S: + """Monoid addition with broadcasting over common collection types, + callables, and interpretations. + + """ + if not args: + return typing.cast(S, self.identity) + + if any(isinstance(x, Term) for x in args): + return typing.cast(S, defdata(self.plus, *args)) + + return self._plus(*args) + + @functools.singledispatchmethod + def _plus[S](self, *args: S) -> S: + return typing.cast(S, functools.reduce(self.kernel, args, self.identity)) + + @_plus.register(Sequence) + def _(self, *args): + return type(args[0])(self.plus(*vs) for vs in zip(*args, strict=True)) + + @_plus.register(Mapping) + def _(self, *args): + if isinstance(args[0], Interpretation): + keys = args[0].keys() + + for b in args[1:]: + if not isinstance(b, Interpretation): + raise TypeError(f"Expected interpretation but got {b}") + + b_keys = b.keys() + if not keys == b_keys: + raise ValueError( + f"Expected interpretation of {keys} but got {b_keys}" + ) + + result = {k: self.plus(*(handler(b)(b[k]) for b in args)) for k in keys} + return result + + for b in args[1:]: + if not isinstance(b, Mapping): + raise TypeError(f"Expected mapping but got {b}") + + all_values = collections.defaultdict(list) + for d in args: + for k, v in d.items(): + all_values[k].append(v) + result = {k: self.plus(*vs) for (k, vs) in all_values.items()} + return result + + @Operation.define + @functools.singledispatchmethod + def reduce[A, B, U: Body]( + self, + body: Annotated[U, Scoped[A | B]], + streams: Annotated[Streams, Scoped[A]], + ) -> Annotated[U, Scoped[B]]: + if callable(body): + return typing.cast(U, lambda *a, **k: self.reduce(body(*a, **k), streams)) + + def generator(loop_order) -> Iterator[Interpretation]: + if len(loop_order) == 0: + return + + stream_key = loop_order[0][0] + stream_values = evaluate(streams[stream_key]) + stream_values_iter = iter(stream_values) # type: ignore[arg-type] + + # If we try to iterate and get a term instead of a real + # iterator, give up + if isinstance(stream_values_iter, Term) and stream_values_iter.op is iter_: + raise NotHandled + + if len(loop_order) == 1: + for val in stream_values_iter: + yield {stream_key: functools.partial(lambda v: v, val)} + else: + for val in stream_values_iter: + intp: Interpretation = { + stream_key: functools.partial(lambda v: v, val) + } + with handler(intp): + for intp2 in generator(loop_order[1:]): + yield coproduct(intp, intp2) + + loop_order = list(order_streams(streams)) + try: + return self.plus( + *(handler(intp)(evaluate)(body) for intp in generator(loop_order)) + ) + except NotHandled: + return typing.cast(U, defdata(self.reduce, body, streams)) + + @reduce.register # type: ignore[attr-defined] + def _(self, body: Mapping, streams): + return {k: self.reduce(v, streams) for (k, v) in body.items()} + + @reduce.register # type: ignore[attr-defined] + def _(self, body: Sequence, streams): + return type(body)(self.reduce(x, streams) for x in body) # type:ignore[call-arg] + + @reduce.register # type: ignore[attr-defined] + def _(self, body: Generator, streams): + return (self.reduce(x, streams) for x in body) + + +class IdempotentMonoid[T](Monoid[T]): + @Operation.define + def plus[S: Body[T]](self, *args: S) -> S: + return super().plus(*args) + + @Operation.define + def reduce[A, B, U: Body]( + self, + body: Annotated[U, Scoped[A | B]], + streams: Annotated[Streams, Scoped[A]], + ) -> Annotated[U, Scoped[B]]: + return super().reduce(body, streams) + + +class CommutativeMonoid[T](Monoid[T]): + @Operation.define + def plus[S: Body[T]](self, *args: S) -> S: + return super().plus(*args) + + @Operation.define + def reduce[A, B, U: Body]( + self, + body: Annotated[U, Scoped[A | B]], + streams: Annotated[Streams, Scoped[A]], + ) -> Annotated[U, Scoped[B]]: + return super().reduce(body, streams) + + +class CommutativeMonoidWithZero[T](CommutativeMonoid[T]): + zero: T + + def __init__(self, kernel: Callable[[T, T], T], identity: T, zero: T): + super().__init__(kernel, identity) + self.zero = zero + + def __repr__(self): + return f"{type(self)}({self.kernel}, {self.identity}, {self.zero})" + + @Operation.define + def plus[S: Body[T]](self, *args: S) -> S: + return super().plus(*args) + + @Operation.define + def reduce[A, B, U: Body]( + self, + body: Annotated[U, Scoped[A | B]], + streams: Annotated[Streams, Scoped[A]], + ) -> Annotated[U, Scoped[B]]: + return super().reduce(body, streams) + + +class Semilattice[T](IdempotentMonoid[T], CommutativeMonoid[T]): + @Operation.define + def plus[S: Body[T]](self, *args: S) -> S: + return super().plus(*args) + + @Operation.define + def reduce[A, B, U: Body]( + self, + body: Annotated[U, Scoped[A | B]], + streams: Annotated[Streams, Scoped[A]], + ) -> Annotated[U, Scoped[B]]: + return super().reduce(body, streams) + + +@Operation.define +def _arg_min[T]( + a: tuple[numbers.Number, T | None], b: tuple[numbers.Number, T | None] +) -> tuple[numbers.Number, T | None]: + if isinstance(a[0], Term) or isinstance(b[0], Term): + raise NotHandled + return b if b[0] < a[0] else a # type: ignore + + +@Operation.define +def _arg_max[T]( + a: tuple[numbers.Number, T | None], b: tuple[numbers.Number, T | None] +) -> tuple[numbers.Number, T | None]: + if isinstance(a[0], Term) or isinstance(b[0], Term): + raise NotHandled + return b if b[0] > a[0] else a # type: ignore + + +Min = Semilattice(kernel=min, identity=float("inf")) +Max = Semilattice(kernel=max, identity=float("-inf")) +ArgMin = Monoid(kernel=_arg_min, identity=(float("inf"), None)) +ArgMax = Monoid(kernel=_arg_max, identity=(float("-inf"), None)) +Sum = CommutativeMonoid(kernel=_NumberTerm.__add__, identity=0) +Product = CommutativeMonoidWithZero(kernel=_NumberTerm.__mul__, identity=1, zero=0) + + +@dataclass +class _ExtensibleBinaryRelation[S, T]: + tuples: set[tuple[S, T]] + + def register(self, s: S, t: T) -> None: + self.tuples.add((s, t)) + + def __call__(self, s: S, t: T) -> bool: + return (s, t) in self.tuples + + +distributes_over = _ExtensibleBinaryRelation( + { + (Max.plus, Min.plus), + (Min.plus, Max.plus), + (Sum.plus, Min.plus), + (Sum.plus, Max.plus), + (Product.plus, Sum.plus), + } +) + + +class PlusEmpty(ObjectInterpretation): + """plus() = 0""" + + @implements(Monoid.plus) + def plus(self, monoid, *args): + if not args: + return monoid.identity + return fwd() + + +class PlusSingle(ObjectInterpretation): + """plus(x) = x""" + + @implements(Monoid.plus) + def plus(self, _, *args): + if len(args) == 1: + return args[0] + return fwd() + + +class PlusIdentity(ObjectInterpretation): + """x₁ + ... + 0 + ... + xₙ = x₁ + ... + xₙ""" + + @implements(Monoid.plus) + def plus(self, monoid, *args): + if any(x is monoid.identity for x in args): + return monoid.plus(*(x for x in args if x is not monoid.identity)) + return fwd() + + +class PlusAssoc(ObjectInterpretation): + """x + (y + z) = (x + y) + z = x + y + z""" + + @implements(Monoid.plus) + def plus(self, monoid, *args): + if any(isinstance(x, Term) and x.op is monoid.plus for x in args): + flat_args = itertools.chain.from_iterable( + t.args if isinstance(t, Term) and t.op is monoid.plus else (t,) + for t in args + ) + assert len(args) > 0 + return monoid.plus(*flat_args) + return fwd() + + +class PlusDistr(ObjectInterpretation): + """x + (y * z) = x * y + x * z""" + + @implements(Monoid.plus) + def plus(self, monoid, *args): + if any( + isinstance(x, Term) and distributes_over(monoid.plus, x.op) for x in args + ): + non_terms = [] + + # group terms by head operation + by_head_op = defaultdict(list) + for t in args: + if isinstance(t, Term): + by_head_op[t.op].append(t) + else: + non_terms.append(t) + + # distribute over each group + progress = False + final_sum = [] + for op, terms in by_head_op.items(): + if ( + len(terms) > 1 + and distributes_over(monoid.plus, op) + and not distributes_over(op, monoid.plus) + ): + progress = True + term_args = (t.args for t in terms) + dist_terms = ( + monoid.plus(*args) for args in itertools.product(*term_args) + ) + final_sum.append(op(*dist_terms)) + else: + final_sum += terms + if progress: + return monoid.plus(*non_terms, *final_sum) + return fwd() + + +class PlusZero(ObjectInterpretation): + """x₁ * ... * 0 * ... * xₙ = 0""" + + @implements(CommutativeMonoidWithZero.plus) + def plus(self, monoid, *args): + if any(x is monoid.zero for x in args): + return monoid.zero + return fwd() + + +class PlusConsecutiveDups(ObjectInterpretation): + """x ⊕ x ⊕ y = x ⊕ y""" + + @implements(IdempotentMonoid.plus) + def plus(self, monoid, *args): + dedup_args = ( + args[i] + for i in range(len(args)) + if i == 0 or not syntactic_eq(args[i - 1], args[i]) + ) + return fwd(monoid, *dedup_args) + + +class PlusDups(ObjectInterpretation): + """x ⊕ y ⊕ x = x ⊕ y""" + + @dataclass + class _HashableTerm: + term: Term + + def __eq__(self, other): + return syntactic_eq(self, other) + + def __hash__(self): + return syntactic_hash(self) + + @implements(Semilattice.plus) + def plus(self, monoid, *args): + # elim dups + args_count = Counter(self._HashableTerm(t) for t in args) + if len(args_count) < len(args): + dedup_args = [] + for t in args: + ht = self._HashableTerm(t) + if ht in args_count: + dedup_args.append(t) + del args_count[ht] + return fwd(monoid, *dedup_args) + return fwd() + + +NormalizePlusIntp = functools.reduce( + coproduct, + typing.cast( + list[Interpretation], + [ + PlusEmpty(), + PlusSingle(), + PlusIdentity(), + PlusAssoc(), + PlusDistr(), + PlusZero(), + PlusConsecutiveDups(), + PlusDups(), + ], + ), +) + + +class ReduceNoStreams(ObjectInterpretation): + """Implements the identity + reduce(R, ∅, body) = 0 + """ + + @implements(Monoid.reduce) + def reduce(self, monoid, _, streams): + if len(streams) == 0: + return monoid.identity + return fwd() + + +class ReduceFusion(ObjectInterpretation): + """Implements the identity + reduce(R, S1, reduce(R, S2, body)) = reduce(R, S1 ∪ S2, body) + """ + + @implements(Monoid.reduce) + def reduce(self, monoid, body, streams): + if isinstance(body, Term) and body.op == monoid.reduce: + return monoid.reduce(body.args[0], streams | body.args[1]) + return fwd() + + +class ReduceSplit(ObjectInterpretation): + """Implements the identity + reduce(R, S, b1 + ... + bn) = reduce(R, S, b1) + ... + reduce(R, S, bn) + """ + + @implements(CommutativeMonoid.reduce) + def reduce(self, monoid, body, streams): + if isinstance(body, Term) and body.op == monoid.plus: + return monoid.plus(*(monoid.reduce(x, streams) for x in body.args)) + return fwd() + + +class ReduceFactorization(ObjectInterpretation): + """ + Implements factorization of independent terms. + For example, when having two independent distributions, + we can rewrite their marginalization as: + ∫p(x)⋅q(y)dxdy => ∫p(x)dx ⋅ ∫q(y)dy + + More specifically, in terms of reduces we are performing: + reduce(R, (S₁ × ... × Sₖ) , A₁ * ... * Aₖ) + => reduce(R, S₁, A₁) * ... * reduce(R, Sₖ, Aₖ) + where free(Aᵢ) ∩ free(Aⱼ) ∩ S = ∅ + and free(Aᵢ) ∩ S ⊆ Sᵢ + """ + + @implements(CommutativeMonoid.reduce) + def reduce(self, monoid, body, streams): + if isinstance(body, Term) and distributes_over(body.op, monoid.plus): + stream_vars = set(streams.keys()) + factors = [(arg, fvsof(arg)) for arg in body.args] + stream_ids = {v: i for (i, v) in enumerate(stream_vars)} + ds = DisjointSet(len(streams)) + + # streams are in the same partition as their dependencies + for stream_var, stream_id in stream_ids.items(): + stream_body = streams[stream_var] + deps = sorted([stream_ids[v] for v in fvsof(stream_body) & stream_vars]) + ds.union(stream_id, *deps) + + # factors are in the same partition as their dependencies + for factor, factor_fvs in factors: + factor_streams = sorted( + [stream_ids[v] for v in (factor_fvs & stream_vars)] + ) + ds.union(*factor_streams) + + placed_streams = set() + new_reduces = [] + for stream_key in streams: + if stream_key in placed_streams: + continue + + partition = ds.find(stream_ids[stream_key]) + partition_streams = { + k: v + for (k, v) in streams.items() + if ds.find(stream_ids[k]) == partition + } + partition_stream_keys = set(partition_streams.keys()) + + partition_factors = [ + t for t in factors if (t[1] & partition_stream_keys) + ] + + assert all( + (t[1] & stream_vars) <= partition_stream_keys + for t in partition_factors + ), "partition contains all streams required by factor" + + partition_term = body.op(*(t[0] for t in partition_factors)) + new_reduces.append((partition_term, partition_streams)) + placed_streams |= partition_stream_keys + + constant_factors = [t for (t, fvs) in factors if not (fvs & stream_vars)] + + if len(new_reduces) > 1: + result = body.op( + *constant_factors, *(monoid.reduce(*args) for args in new_reduces) + ) + return result + + return fwd() + + +NormalizeReduceIntp = functools.reduce( + coproduct, + typing.cast( + list[Interpretation], + [ReduceNoStreams(), ReduceFusion(), ReduceSplit(), ReduceFactorization()], + ), +) + +NormalizeIntp = coproduct(NormalizePlusIntp, NormalizeReduceIntp) diff --git a/effectful/ops/syntax.py b/effectful/ops/syntax.py index d73546028..a2fbcf9b2 100644 --- a/effectful/ops/syntax.py +++ b/effectful/ops/syntax.py @@ -903,6 +903,84 @@ def _(x: object, other) -> bool: return x == other +@_CustomSingleDispatchCallable +def syntactic_hash(__dispatch: Callable[[type], Callable[[Any], int]], x) -> int: + """Structural hash compatible with :func:`syntactic_eq`. + + Guarantees that ``syntactic_eq(x, y)`` implies + ``syntactic_hash(x) == syntactic_hash(y)``. + + :param x: A term. + :returns: An integer hash. + """ + if dataclasses.is_dataclass(x) and not isinstance(x, type): + return hash( + ( + "dataclass", + type(x), + syntactic_hash( + { + field.name: getattr(x, field.name) + for field in dataclasses.fields(x) + } + ), + ) + ) + else: + return __dispatch(type(x))(x) + + +@syntactic_hash.register +def _(x: Term) -> int: + return hash( + ( + "term", + x.op, + len(x.args), + tuple(syntactic_hash(a) for a in x.args), + # sort kwargs so order doesn't affect the hash + tuple((k, syntactic_hash(x.kwargs[k])) for k in sorted(x.kwargs)), + ) + ) + + +@syntactic_hash.register +def _(x: collections.abc.Mapping) -> int: + # XOR over (key_hash, value_hash) pairs — order-independent, + # matching the set-based comparison in syntactic_eq's Mapping branch. + acc = 0 + for k in x: + acc ^= hash((hash(k), syntactic_hash(x[k]))) + return hash(("mapping", acc)) + + +@syntactic_hash.register +def _(x: collections.abc.Sequence) -> int: + if ( + isinstance(x, tuple) + and hasattr(x, "_fields") + and all(hasattr(x, f) for f in x._fields) + ): + return hash( + ( + "namedtuple", + type(x), + tuple(syntactic_hash(getattr(x, f)) for f in x._fields), + ) + ) + else: + # Use the abstract Sequence tag (not type(x)) because syntactic_eq + # treats any two Sequences of equal length and elementwise-equal + # contents as equal — e.g. [1,2] and (1,2) compare equal. + return hash(("sequence", len(x), tuple(syntactic_hash(a) for a in x))) + + +@syntactic_hash.register(object) +@syntactic_hash.register(str | bytes) +def _(x: object) -> int: + return hash(x) + + class ObjectInterpretation[T, V](collections.abc.Mapping): """A helper superclass for defining an ``Interpretation`` of many :class:`~effectful.ops.types.Operation` instances with shared state or behavior. diff --git a/pyproject.toml b/pyproject.toml index 0ae2f45b6..054763c27 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -73,6 +73,7 @@ test = [ "pytest-cov", "pytest-xdist", "pytest-benchmark", + "hypothesis", "mypy", "ruff", "nbval", diff --git a/tests/_monoid_helpers.py b/tests/_monoid_helpers.py new file mode 100644 index 000000000..4532ae72d --- /dev/null +++ b/tests/_monoid_helpers.py @@ -0,0 +1,85 @@ +from collections.abc import Callable, Mapping, Sequence +from typing import Any, get_args, get_origin + +from hypothesis import strategies as st + +from effectful.ops.syntax import deffn +from effectful.ops.types import Operation + + +def _value_strategy_for(annotation: Any) -> st.SearchStrategy[Any]: + """Strategy for the value an *0-arg* Operation should return.""" + if annotation is int: + return st.integers() + if annotation is float: + return st.floats(allow_nan=False) + if get_origin(annotation) is list and get_args(annotation) == (int,): + return st.lists(st.integers()) + raise NotImplementedError( + f"No value strategy for return annotation {annotation!r}; " + "supported: int, list[int]" + ) + + +_UNARY_INT_FNS: list[Callable[[int], int]] = [ + lambda x: x, + lambda x: x + 1, + lambda x: x - 1, + lambda x: -x, + lambda x: 2 * x, + lambda x: 3 * x + 1, +] + +_BINARY_INT_FNS: list[Callable[[int, int], int]] = [ + lambda x, y: x + y, + lambda x, y: x - y, + lambda x, y: x * y, + lambda x, y: x + 2 * y, + lambda x, y: 2 * x - y, +] + +_UNARY_LIST_FNS: list[Callable[[int], list[int]]] = [ + lambda _x: [], + lambda x: [x], + lambda x: [x, x + 1], + lambda x: [x, -x], + lambda x: [0, x, x + 1], +] + + +def _strategy_for_op(op: Operation) -> st.SearchStrategy[Callable[..., Any]]: + """Pick a strategy producing a callable suitable for binding `op` in an + interpretation. Inspects the operation's signature. + """ + sig = op.__signature__ + params = list(sig.parameters.values()) + ret = sig.return_annotation + param_types = tuple(p.annotation for p in params) + + if not params: + return _value_strategy_for(ret).map(deffn) + if ret is int and param_types == (int,): + return st.sampled_from(_UNARY_INT_FNS) + if ret is int and param_types == (int, int): + return st.sampled_from(_BINARY_INT_FNS) + if get_origin(ret) is list and get_args(ret) == (int,) and param_types == (int,): + return st.sampled_from(_UNARY_LIST_FNS) + raise NotImplementedError( + f"Function-typed free var must return int or list[int]; got {ret!r} for {op}" + ) + + +@st.composite +def random_interpretation( + draw: st.DrawFn, free_vars: Sequence[Operation] +) -> Mapping[Operation, Callable[..., Any]]: + """Draw an Interpretation binding every Operation in `case.free_vars` to + a randomly chosen value/callable. Keys are Operation identities. + """ + intp: dict[Operation, Callable[..., Any]] = {} + for op in free_vars: + intp[op] = draw(_strategy_for_op(op)) + return intp + + +__all__ = ["random_interpretation"] diff --git a/tests/test_internals_disjoint_set.py b/tests/test_internals_disjoint_set.py new file mode 100644 index 000000000..808b8d25d --- /dev/null +++ b/tests/test_internals_disjoint_set.py @@ -0,0 +1,124 @@ +import random + +import pytest + +from effectful.internals.disjoint_set import DisjointSet + + +@pytest.fixture +def dsu(): + return DisjointSet(10) + + +def test_initial_state(dsu): + for i in range(10): + assert dsu.find(i) == i + + +def test_simple_union(dsu): + assert dsu.union(1, 2) is True + assert dsu.find(1) == dsu.find(2) + + +def test_union_idempotent(dsu): + dsu.union(1, 2) + assert dsu.union(1, 2) is False + + +def test_union_chain(dsu): + dsu.union(1, 2) + dsu.union(2, 3) + assert dsu.find(1) == dsu.find(3) + + +def test_union_multiple_elements_all_connected(dsu): + dsu.union(1, 2, 3, 4, 5) + roots = {dsu.find(i) for i in [1, 2, 3, 4, 5]} + assert len(roots) == 1 + + +def test_union_multiple_elements_partial_overlap(dsu): + dsu.union(1, 2) + dsu.union(3, 4) + dsu.union(2, 3, 5) + + roots = {dsu.find(i) for i in [1, 2, 3, 4, 5]} + assert len(roots) == 1 + + +def test_union_multiple_elements_with_existing_connections(dsu): + dsu.union(1, 2) + dsu.union(2, 3) + dsu.union(3, 4, 5, 6) + + roots = {dsu.find(i) for i in [1, 2, 3, 4, 5, 6]} + assert len(roots) == 1 + + +def test_union_single_element(dsu): + assert dsu.union(1) is False + + +def test_union_no_elements(dsu): + assert dsu.union() is False + + +def test_union_self(dsu): + assert dsu.union(3, 3) is False + assert dsu.find(3) == 3 + + +def test_transitivity(dsu): + dsu.union(1, 2) + dsu.union(2, 3) + dsu.union(3, 4) + assert dsu.find(1) == dsu.find(4) + + +def test_disjoint_sets_remain_separate(dsu): + dsu.union(1, 2) + dsu.union(3, 4) + assert dsu.find(1) != dsu.find(3) + + +def test_randomized_unions(): + n = 50 + dsu = DisjointSet(n) + + groups = [{i} for i in range(n)] + + def find_group(x): + for g in groups: + if x in g: + return g + + for _ in range(100): + elems = random.sample(range(n), random.randint(2, 5)) + dsu.union(*elems) + + # merge ground-truth groups + merged = set() + for e in elems: + merged |= find_group(e) + + groups = [g for g in groups if g.isdisjoint(merged)] + groups.append(merged) + + # verify structure matches ground truth + for g in groups: + roots = {dsu.find(x) for x in g} + assert len(roots) == 1 + + +def test_path_compression_effect(): + dsu = DisjointSet(6) + dsu.union(0, 1) + dsu.union(1, 2) + dsu.union(2, 3) + dsu.union(3, 4) + + # Trigger compression + root_before = dsu.find(4) + root_after = dsu.find(4) + + assert root_before == root_after diff --git a/tests/test_ops_monoid.py b/tests/test_ops_monoid.py new file mode 100644 index 000000000..a22928cca --- /dev/null +++ b/tests/test_ops_monoid.py @@ -0,0 +1,518 @@ +import functools +import itertools + +import pytest +from hypothesis import given, settings +from hypothesis import strategies as st + +from effectful.internals.runtime import interpreter +from effectful.ops.monoid import Max, Min, NormalizeIntp, Product, Semilattice, Sum +from effectful.ops.semantics import apply, evaluate, fvsof, handler +from effectful.ops.syntax import _BaseTerm, defdata, syntactic_eq +from effectful.ops.types import NotHandled, Operation +from tests._monoid_helpers import random_interpretation + +_INT = st.integers(min_value=-100, max_value=100) + +ALL_MONOIDS = [ + pytest.param(Sum, id="Sum"), + pytest.param(Product, id="Product"), + pytest.param(Min, id="Min"), + pytest.param(Max, id="Max"), +] + +COMMUTATIVE = [ + pytest.param(Sum, id="Sum"), + pytest.param(Product, id="Product"), + pytest.param(Min, id="Min"), + pytest.param(Max, id="Max"), +] + +IDEMPOTENT = [ + pytest.param(Min, id="Min"), + pytest.param(Max, id="Max"), +] + +WITH_ZERO = [ + pytest.param(Product, id="Product"), +] + + +def define_vars(*names, typ=int): + if len(names) == 1: + return Operation.define(typ, name=names[0]) + return tuple(Operation.define(typ, name=n) for n in names) + + +@functools.cache +def _canonical_op(idx: int) -> Operation: + """Globally cached canonical Operation, keyed by encounter index. + + Cached so that two independent canonicalize runs return the same + Operation object for the same index — letting ``syntactic_eq`` + compare canonical forms by Operation identity. + """ + return Operation.define(int, name=f"__cv_{idx}") + + +def syntactic_eq_alpha(x, y) -> bool: + """Alpha-equivalence-respecting variant of ``syntactic_eq``. + + Walks each expression bottom-up with :func:`evaluate` and renames + every bound variable to a deterministic canonical Operation. The + canonical names are assigned by a counter that increments in + ``evaluate``'s natural traversal order, so two alpha-equivalent + expressions canonicalize to syntactically identical results. + """ + return syntactic_eq(_canonicalize(x), _canonicalize(y)) + + +def _canonicalize(expr): + counter = itertools.count() + + def _passthrough(op, *args, **kwargs): + return defdata(op, *args, **kwargs) + + def _substitute(arg, renaming): + """Apply a bound-variable renaming using ``evaluate`` for traversal.""" + if not renaming: + return arg + with interpreter({apply: _passthrough, **renaming}): + return evaluate(arg) + + def _bound_var_order(args, kwargs, bound_set): + """Return bound variables in deterministic encounter order.""" + seen: list[Operation] = [] + seen_set: set[Operation] = set() + + def _capture(op, *a, **kw): + if op in bound_set and op not in seen_set: + seen.append(op) + seen_set.add(op) + return defdata(op, *a, **kw) + + # ``evaluate`` walks Terms, lists, tuples, mappings, dataclasses, + # etc. for free; the apply handler captures bound vars used as + # ``x()`` anywhere in the body. + with interpreter({apply: _capture}): + evaluate((args, kwargs)) + + # Binders bypass the apply handler. Pick them up with a small structural + # walk that visits dict keys too. + def _walk_bare(obj): + if isinstance(obj, Operation): + if obj in bound_set and obj not in seen_set: + seen.append(obj) + seen_set.add(obj) + elif isinstance(obj, dict): + for k, v in obj.items(): + _walk_bare(k) + _walk_bare(v) + elif isinstance(obj, list | set | frozenset | tuple): + for v in obj: + _walk_bare(v) + + _walk_bare((args, kwargs)) + return seen + + def _apply_canonical(op, *args, **kwargs): + bindings = op.__fvs_rule__(*args, **kwargs) + all_bound: set[Operation] = set().union( + *bindings.args, *bindings.kwargs.values() + ) + if not all_bound: + return defdata(op, *args, **kwargs) + + order = _bound_var_order(args, kwargs, all_bound) + canonical = {var: _canonical_op(next(counter)) for var in order} + assert all_bound <= set(order) + + new_args = tuple( + _substitute( + arg, {v: canonical[v] for v in bindings.args[i] if v in canonical} + ) + for i, arg in enumerate(args) + ) + new_kwargs = { + k: _substitute( + v, + {var: canonical[var] for var in bindings.kwargs[k] if var in canonical}, + ) + for k, v in kwargs.items() + } + + # avoid the renaming from defdata + return _BaseTerm(op, *new_args, **new_kwargs) + + with interpreter({apply: _apply_canonical}): + return evaluate(expr) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +@given(a=_INT, b=_INT, c=_INT) +@settings(max_examples=50, deadline=None) +def test_associativity(monoid, a, b, c): + left = monoid.plus(monoid.plus(a, b), c) + right = monoid.plus(a, monoid.plus(b, c)) + assert left == right + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +@given(a=_INT) +@settings(max_examples=50, deadline=None) +def test_identity(monoid, a): + assert monoid.plus(monoid.identity, a) == a + assert monoid.plus(a, monoid.identity) == a + + +@pytest.mark.parametrize("monoid", COMMUTATIVE) +@given(a=_INT, b=_INT) +@settings(max_examples=50, deadline=None) +def test_commutativity(monoid, a, b): + assert monoid.plus(a, b) == monoid.plus(b, a) + + +@pytest.mark.parametrize("monoid", IDEMPOTENT) +@given(a=_INT) +@settings(max_examples=50, deadline=None) +def test_idempotence(monoid, a): + assert monoid.plus(a, a) == a + + +@pytest.mark.parametrize("monoid", WITH_ZERO) +@given(a=_INT) +@settings(max_examples=50, deadline=None) +def test_zero_absorbs(monoid, a): + assert monoid.plus(monoid.zero, a) == monoid.zero + assert monoid.plus(a, monoid.zero) == monoid.zero + + +def _check_pair(lhs, rhs, *, free_vars=[], max_examples: int = 25) -> None: + """Run structural + semantic checks on a TermPair.""" + with handler(NormalizeIntp): + norm = evaluate(lhs) + + assert syntactic_eq_alpha(norm, rhs) + + @given(intp=random_interpretation(free_vars)) + @settings(max_examples=max_examples, deadline=None) + def _check_semantics(intp): + with handler(intp): + lhs_val = evaluate(lhs) + rhs_val = evaluate(rhs) + assert lhs_val == rhs_val + + _check_semantics() + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_plus_empty(monoid): + _check_pair(lhs=monoid.plus(), rhs=monoid.identity) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_plus_single(monoid): + x = define_vars("x", typ=type(monoid.identity)) + _check_pair(lhs=monoid.plus(x()), rhs=x(), free_vars=[x]) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_plus_identity_right(monoid): + x = define_vars("x", typ=type(monoid.identity)) + _check_pair(lhs=monoid.plus(x(), monoid.identity), rhs=x(), free_vars=[x]) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_plus_identity_left(monoid): + x = define_vars("x", typ=type(monoid.identity)) + _check_pair(lhs=monoid.plus(monoid.identity, x()), rhs=x(), free_vars=[x]) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_plus_assoc_right(monoid): + x, y, z = define_vars("x", "y", "z", typ=type(monoid.identity)) + _check_pair( + lhs=monoid.plus(x(), monoid.plus(y(), z())), + rhs=monoid.plus(x(), y(), z()), + free_vars=[x, y, z], + ) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_plus_assoc_left(monoid): + x, y, z = define_vars("x", "y", "z", typ=type(monoid.identity)) + _check_pair( + lhs=monoid.plus(monoid.plus(x(), y()), z()), + rhs=monoid.plus(x(), y(), z()), + free_vars=[x, y, z], + ) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_plus_sequence(monoid): + a, b, c, d = define_vars("a", "b", "c", "d", typ=type(monoid.identity)) + _check_pair( + lhs=monoid.plus([a(), b()], [c(), d()]), + rhs=[monoid.plus(a(), c()), monoid.plus(b(), d())], + free_vars=[a, b, c, d], + ) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_plus_mapping(monoid): + a, b, c, d = define_vars("a", "b", "c", "d", typ=type(monoid.identity)) + _check_pair( + lhs=monoid.plus({"x": a(), "y": b()}, {"x": c(), "z": d()}), + rhs={"x": monoid.plus(a(), c()), "y": b(), "z": d()}, + free_vars=[a, b, c, d], + ) + + +def test_plus_distributes(): + a, b, c, d = define_vars("a", "b", "c", "d") + lhs = Product.plus(Sum.plus(a(), b()), Sum.plus(c(), d())) + rhs = Sum.plus( + Product.plus(a(), c()), + Product.plus(a(), d()), + Product.plus(b(), c()), + Product.plus(b(), d()), + ) + _check_pair(lhs=lhs, rhs=rhs, free_vars=[a, b, c, d]) + + +def test_plus_distributes_constant(): + a, b, c, d = define_vars("a", "b", "c", "d") + lhs = Product.plus(Sum.plus(a(), b()), Sum.plus(c(), d()), 5) + rhs = Product.plus( + 5, + Sum.plus( + Product.plus(a(), c()), + Product.plus(a(), d()), + Product.plus(b(), c()), + Product.plus(b(), d()), + ), + ) + _check_pair(lhs=lhs, rhs=rhs, free_vars=[a, b, c, d]) + + +def test_plus_distributes_multiple(): + a, b, c, d = define_vars("a", "b", "c", "d") + lhs = Sum.plus( + Min.plus(a(), b()), + Min.plus(c(), d()), + Max.plus(a(), b()), + Max.plus(c(), d()), + ) + rhs = Sum.plus( + Min.plus( + Sum.plus(a(), c()), + Sum.plus(a(), d()), + Sum.plus(b(), c()), + Sum.plus(b(), d()), + ), + Max.plus( + Sum.plus(a(), c()), + Sum.plus(a(), d()), + Sum.plus(b(), c()), + Sum.plus(b(), d()), + ), + ) + _check_pair(lhs=lhs, rhs=rhs, free_vars=[a, b, c, d]) + + +@pytest.mark.parametrize("monoid", IDEMPOTENT) +def test_plus_idempotent_consecutive(monoid): + """``a, a, b → a, b`` — only consecutive duplicates collapse.""" + a, b = define_vars("a", "b") + lhs = monoid.plus(a(), a(), b()) + return _check_pair(lhs=lhs, rhs=monoid.plus(a(), b()), free_vars=[a, b]) + + +@pytest.mark.parametrize("monoid", IDEMPOTENT) +def test_plus_idempotent_non_consecutive(monoid): + """``a, b, a`` — Semilattice (Min/Max) collapses via commutative + PlusDups; plain IdempotentMonoid leaves it as-is (consecutive-only).""" + a, b = define_vars("a", "b") + lhs = monoid.plus(a(), b(), a()) + if isinstance(monoid, Semilattice): + rhs = monoid.plus(a(), b()) + else: + rhs = monoid.plus(a(), b(), a()) + _check_pair(lhs=lhs, rhs=rhs, free_vars=[a, b]) + + +def test_plus_commutative_idempotent_long(): + """Long alternation collapses via commutative dedup (Min/Max only).""" + a, b = define_vars("a", "b") + lhs = Min.plus(a(), b(), a(), b(), b(), a(), a()) + _check_pair(lhs=lhs, rhs=Min.plus(a(), b()), free_vars=[a, b]) + + +@pytest.mark.parametrize("monoid", WITH_ZERO) +def test_plus_zero(monoid): + a = define_vars("a") + lhs_right = monoid.plus(a(), monoid.zero) + lhs_left = monoid.plus(monoid.zero, a()) + _check_pair(lhs=lhs_right, rhs=monoid.zero, free_vars=[a]) + _check_pair(lhs=lhs_left, rhs=monoid.zero, free_vars=[a]) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_reduce_body_sequence(monoid): + x = Operation.define(int, name="x") + X = Operation.define(list[int], name="X") + + @Operation.define + def f(_x: int) -> int: + raise NotHandled + + g = Operation.define(f, name="g") + + lhs = monoid.reduce([f(x()), g(x())], {x: X()}) + rhs = [monoid.reduce(f(x()), {x: X()}), monoid.reduce(g(x()), {x: X()})] + + _check_pair(lhs=lhs, rhs=rhs, free_vars=[X, f, g]) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_reduce_body_sequence_2(monoid): + x, y = define_vars("x", "y") + X, Y = define_vars("X", "Y", typ=list[int]) + + @Operation.define + def f(_x: int) -> int: + raise NotHandled + + g = Operation.define(f, name="g") + + lhs = monoid.reduce([f(x()), g(y())], {x: X(), y: Y()}) + rhs = [ + monoid.reduce(f(x()), {x: X(), y: Y()}), + monoid.reduce(g(y()), {x: X(), y: Y()}), + ] + + _check_pair(lhs=lhs, rhs=rhs, free_vars=[X, Y, f, g]) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_reduce_body_mapping(monoid): + x = Operation.define(int, name="x") + X = Operation.define(list[int], name="X") + + @Operation.define + def f(_x: int) -> int: + raise NotHandled + + g = Operation.define(f, name="g") + + lhs = monoid.reduce({"a": f(x()), "b": g(x())}, {x: X()}) + rhs = { + "a": monoid.reduce(f(x()), {x: X()}), + "b": monoid.reduce(g(x()), {x: X()}), + } + _check_pair(lhs=lhs, rhs=rhs, free_vars=[X, f, g]) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_reduce_no_streams(monoid): + a = define_vars("a") + lhs = monoid.reduce(a(), {}) + rhs = monoid.identity + + _check_pair(lhs=lhs, rhs=rhs, free_vars=[a]) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_reduce_reduce(monoid): + a, b = define_vars("a", "b") + A, B = define_vars("A", "B", typ=list[int]) + + @Operation.define + def f(_x: int, _y: int) -> int: + raise NotHandled + + lhs = monoid.reduce(monoid.reduce(f(a(), b()), {a: A()}), {b: B()}) + rhs = monoid.reduce(f(a(), b()), {a: A(), b: B()}) + + _check_pair(lhs=lhs, rhs=rhs, free_vars=[A, B, f]) + + +@pytest.mark.parametrize("monoid", COMMUTATIVE) +def test_reduce_plus(monoid): + a, b = define_vars("a", "b") + A, B = define_vars("A", "B", typ=list[int]) + lhs = monoid.reduce(monoid.plus(a(), b()), {a: A(), b: B()}) + rhs = monoid.plus( + monoid.reduce(a(), {a: A(), b: B()}), + monoid.reduce(b(), {a: A(), b: B()}), + ) + _check_pair(lhs=lhs, rhs=rhs, free_vars=[A, B]) + + +def test_reduce_independent_1(): + a, b = define_vars("a", "b") + A, B = define_vars("A", "B", typ=list[int]) + lhs = Sum.reduce(Product.plus(a(), b()), {a: A(), b: B()}) + rhs = Product.plus(Sum.reduce(a(), {a: A()}), Sum.reduce(b(), {b: B()})) + _check_pair(lhs=lhs, rhs=rhs, free_vars=[A, B]) + + +def test_reduce_independent_2(): + a, b, c = define_vars("a", "b", "c") + A, B, C = define_vars("A", "B", "C", typ=list[int]) + + @Operation.define + def f(_x: int, _y: int) -> int: + raise NotHandled + + lhs = Sum.reduce(Product.plus(a(), b(), f(b(), c())), {a: A(), b: B(), c: C()}) + rhs = Product.plus( + Sum.reduce(a(), {a: A()}), + Sum.reduce(Product.plus(b(), f(b(), c())), {b: B(), c: C()}), + ) + _check_pair(lhs=lhs, rhs=rhs, free_vars=[A, B, C, f]) + + +def test_reduce_independent_3_negative(): + """Stream `b` depends on `a` (b: g(a())), so the proposed factorization + is unsound — the normalizer must NOT apply it.""" + a, b, c = define_vars("a", "b", "c") + A, C = define_vars("A", "C", typ=list[int]) + + @Operation.define + def f(_x: int, _y: int) -> int: + raise NotHandled + + @Operation.define + def g(_x: int) -> list[int]: + raise NotHandled + + with handler(NormalizeIntp): + lhs = Sum.reduce( + Product.plus(a(), b(), f(b(), c())), {a: A(), b: g(a()), c: C()} + ) + bogus_rhs = Product.plus( + Sum.reduce(a(), {a: A()}), + Sum.reduce(Product.plus(b(), f(b(), c())), {b: g(a()), c: C()}), + ) + assert fvsof(bogus_rhs) != fvsof(lhs) + # Structural-only negative check: the normalizer correctly refused to apply + # the bogus factorization. + assert not syntactic_eq_alpha(lhs, bogus_rhs) + + +def test_reduce_independent_4(): + a, b, c = define_vars("a", "b", "c") + A, B, C = define_vars("A", "B", "C", typ=list[int]) + + @Operation.define + def f(_x: int, _y: int) -> int: + raise NotHandled + + lhs = Sum.reduce(Product.plus(a(), b(), f(b(), c()), 7), {a: A(), b: B(), c: C()}) + rhs = Product.plus( + 7, + Sum.reduce(a(), {a: A()}), + Sum.reduce(Product.plus(b(), f(b(), c())), {b: B(), c: C()}), + ) + _check_pair(lhs=lhs, rhs=rhs, free_vars=[A, B, C, f]) From fa85541b31e42d408ae208571d0d200d66c30b43 Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Tue, 12 May 2026 14:57:58 -0400 Subject: [PATCH 02/16] Add inversion from `weighted` (#655) * Add monoid module (#653) * add monoid module * clean up * fix doctest * fix * wip * remove incorrect rule * add disjoint set tests and fix bug * lint * drop jax monoid defs * drop incorrect comment * add assert * reduce nondeterminism and add assertions * fix inconsistent stream numbering and missing constant factors * wip * cleanup * fix rule * wip * fix bug * cleanup * lin --- effectful/internals/product_n.py | 4 +- effectful/ops/monoid.py | 162 ++++++++++++++++++++++++++-- effectful/ops/semantics.py | 1 + effectful/ops/types.py | 5 +- tests/_monoid_helpers.py | 14 +-- tests/test_handlers_llm_provider.py | 2 +- tests/test_ops_monoid.py | 146 ++++++++++++++++++++++--- tests/test_ops_syntax.py | 1 - 8 files changed, 301 insertions(+), 34 deletions(-) diff --git a/effectful/internals/product_n.py b/effectful/internals/product_n.py index ca7c2f0b3..e96a2f33a 100644 --- a/effectful/internals/product_n.py +++ b/effectful/internals/product_n.py @@ -79,8 +79,8 @@ def recurse(x): else: result = type(expr)(recurse(tuple(expr.items()))) elif isinstance(expr, collections.abc.Sequence): - if isinstance(expr, str | bytes): - result = expr + if isinstance(expr, str | bytes | range): + return expr elif ( isinstance(expr, tuple) and hasattr(expr, "_fields") diff --git a/effectful/ops/monoid.py b/effectful/ops/monoid.py index 58a10ba3d..ad83de47b 100644 --- a/effectful/ops/monoid.py +++ b/effectful/ops/monoid.py @@ -4,24 +4,26 @@ import numbers import typing from collections import Counter, defaultdict -from collections.abc import Callable, Generator, Iterable, Iterator, Mapping, Sequence +from collections.abc import Callable, Generator, Iterable, Iterator, Mapping from dataclasses import dataclass from graphlib import TopologicalSorter from typing import Annotated, Any from effectful.internals.disjoint_set import DisjointSet +from effectful.internals.runtime import interpreter from effectful.ops.semantics import coproduct, evaluate, fvsof, fwd, handler from effectful.ops.syntax import ( ObjectInterpretation, Scoped, _NumberTerm, defdata, + deffn, implements, iter_, syntactic_eq, syntactic_hash, ) -from effectful.ops.types import Interpretation, NotHandled, Operation, Term +from effectful.ops.types import Expr, Interpretation, NotHandled, Operation, Term # Note: The streams value type should be something like Iterable[T], but some of # our target stream types (e.g. jax.Array) are not subtypes of Iterable @@ -80,9 +82,13 @@ def plus[S: Body[T]](self, *args: S) -> S: def _plus[S](self, *args: S) -> S: return typing.cast(S, functools.reduce(self.kernel, args, self.identity)) - @_plus.register(Sequence) + @_plus.register(tuple) def _(self, *args): - return type(args[0])(self.plus(*vs) for vs in zip(*args, strict=True)) + return tuple(self.plus(*vs) for vs in zip(*args, strict=True)) + + @_plus.register(Generator) + def _(self, *args): + return (self.plus(*vs) for vs in zip(*args, strict=True)) @_plus.register(Mapping) def _(self, *args): @@ -161,8 +167,8 @@ def _(self, body: Mapping, streams): return {k: self.reduce(v, streams) for (k, v) in body.items()} @reduce.register # type: ignore[attr-defined] - def _(self, body: Sequence, streams): - return type(body)(self.reduce(x, streams) for x in body) # type:ignore[call-arg] + def _(self, body: tuple, streams): + return tuple(self.reduce(x, streams) for x in body) @reduce.register # type: ignore[attr-defined] def _(self, body: Generator, streams): @@ -252,12 +258,26 @@ def _arg_max[T]( return b if b[0] > a[0] else a # type: ignore +@Operation.define +def product[T]( + a: Iterable[tuple[T, ...] | T], b: Iterable[tuple[T, ...] | T] +) -> Iterable[tuple[T, ...]]: + if isinstance(a, Term) or isinstance(b, Term): + raise NotHandled + + def to_tuple(x): + return x if isinstance(x, tuple) else (x,) + + return [to_tuple(x) + to_tuple(y) for (x, y) in itertools.product(a, b)] + + Min = Semilattice(kernel=min, identity=float("inf")) Max = Semilattice(kernel=max, identity=float("-inf")) ArgMin = Monoid(kernel=_arg_min, identity=(float("inf"), None)) ArgMax = Monoid(kernel=_arg_max, identity=(float("-inf"), None)) Sum = CommutativeMonoid(kernel=_NumberTerm.__add__, identity=0) Product = CommutativeMonoidWithZero(kernel=_NumberTerm.__mul__, identity=1, zero=0) +CartesianProduct = Monoid(kernel=product, identity=[()]) @dataclass @@ -545,11 +565,139 @@ def reduce(self, monoid, body, streams): return fwd() +def inner_stream( + streams: dict[Operation, Expr], +) -> Iterable[tuple[dict[Operation, Expr], Operation, Expr]]: + """Returns the streams that can be ordered innermost in the loop nest as + well as the remaining streams in the nest. + + """ + stream_vars = set(streams.keys()) + + no_dependents = set() + succ = defaultdict(set) + for k, v in streams.items(): + preds = fvsof(v) & stream_vars + if preds: + for pred in preds: + succ[pred].add(k) + else: + no_dependents.add(k) + + topo = TopologicalSorter(succ) + topo.prepare() + return ( + ({k: v for (k, v) in streams.items() if k != op}, op, streams[op]) + for op in set(topo.get_ready()) | no_dependents + ) + + +def match_reduce(term: Term) -> tuple | None: + reduce_args = None + + def set_reduce_args(*args, **kwargs): + nonlocal reduce_args + reduce_args = args + + with interpreter({Monoid.reduce: set_reduce_args}): + term.op(*term.args, **term.kwargs) + return reduce_args + + +class ReduceDistributeCartesianProduct(ObjectInterpretation): + """Eliminates a reduce over a cartesian product. + ∑_x₁ ∑_x₂ ... ∑_xₙ ∏_i f(xᵢ) = ∏_i ∑_xᵢ f(xᵢ) + This transform is also called inversion in the lifting + literature (e.g. [1]). + + More specifically, this transform implements the identity + reduce(⨁, reduce(⨂, body2, {vv: v()}), {v: reduce(×, body1, S1)} ∪ S2) + = reduce(⨁, reduce(⨂, reduce(⨁, body2, {vv: body1}), S1), S2) + where × is the cartesian product and ⨂ distributes over ⨁. + + Note: This could be generalized to grouped inversion [2]. + + [1] Braz, Rd, Eyal Amir, and Dan Roth. "Lifted first-order + probabilistic inference." IJCAI. 2005. + [2] Taghipour, Nima, et al. "Completeness results for lifted + variable elimination." AISTATS. 2013. + """ + + @implements(CommutativeMonoid.reduce) + def reduce(self, sum_monoid: Monoid, sum_body, sum_streams): + if not (isinstance(sum_body, Term)): + return fwd() + + # body is a product or multiplication of products + if distributes_over(sum_body.op, sum_monoid.plus): + prod_reduces = sum_body.args + else: + prod_reduces = [sum_body] + + products: list[tuple[Monoid, Callable, Operation, Term]] = [] + for prod_reduce in prod_reduces: + prod_args = match_reduce(prod_reduce) + if prod_args is None: + return fwd() + (prod_monoid, prod_body, prod_streams) = prod_args + if not ( + distributes_over(prod_monoid.plus, sum_monoid.plus) + and (len(products) == 0 or products[-1][0] == prod_monoid) + ): + return fwd() + + if len(prod_streams) > 1 or len(prod_streams) == 0: + return fwd() + (prod_op, prod_stream) = next(iter(prod_streams.items())) + products.append( + (prod_monoid, deffn(prod_body, prod_op), prod_op, prod_stream) + ) + + assert len(products) > 0 + + for outer_sum_streams, cprod_op, cprod_term in inner_stream(sum_streams): + if not ( + isinstance(cprod_term, Term) + and cprod_term.op == CartesianProduct.reduce + ): + continue + (cprod_body, cprod_streams) = cprod_term.args + + if not all( + prod_stream.op == cprod_op for (_, _, _, prod_stream) in products + ): + continue + + prod_op = Operation.define(products[0][2]) + prod_monoid = products[0][0] + inner_sum = sum_monoid.reduce( + prod_monoid.plus( + *(prod_body(prod_op()) for (_, prod_body, _, _) in products) + ), + {prod_op: cprod_body}, + ) + prod = prod_monoid.reduce(inner_sum, cprod_streams) + outer_sum = ( + sum_monoid.reduce(prod, outer_sum_streams) + if outer_sum_streams + else prod + ) + return outer_sum + + return fwd() + + NormalizeReduceIntp = functools.reduce( coproduct, typing.cast( list[Interpretation], - [ReduceNoStreams(), ReduceFusion(), ReduceSplit(), ReduceFactorization()], + [ + ReduceNoStreams(), + ReduceFusion(), + ReduceSplit(), + ReduceFactorization(), + ReduceDistributeCartesianProduct(), + ], ), ) diff --git a/effectful/ops/semantics.py b/effectful/ops/semantics.py index 4fd5c481c..16ace787a 100644 --- a/effectful/ops/semantics.py +++ b/effectful/ops/semantics.py @@ -173,6 +173,7 @@ def evaluate[T]( @evaluate.register(object) @evaluate.register(str) @evaluate.register(bytes) +@evaluate.register(range) def _evaluate_object[T](expr: T, **kwargs) -> T: if dataclasses.is_dataclass(expr) and not isinstance(expr, type): return _evaluate_dataclass(expr, **kwargs) diff --git a/effectful/ops/types.py b/effectful/ops/types.py index 96cda80ed..34379554e 100644 --- a/effectful/ops/types.py +++ b/effectful/ops/types.py @@ -496,7 +496,10 @@ def _instance_op(instance, *args, **kwargs): else: return default_result - instance_op = self.define(types.MethodType(_instance_op, instance)) + name = ("" if owner is None else f"{owner.__name__}_") + self.__name__ + instance_op = self.define( + types.MethodType(_instance_op, instance), name=name + ) instance.__dict__[self._name_on_instance] = instance_op return instance_op elif instance is not None: diff --git a/tests/_monoid_helpers.py b/tests/_monoid_helpers.py index 4532ae72d..9b311b257 100644 --- a/tests/_monoid_helpers.py +++ b/tests/_monoid_helpers.py @@ -14,14 +14,14 @@ def _value_strategy_for(annotation: Any) -> st.SearchStrategy[Any]: if annotation is float: return st.floats(allow_nan=False) if get_origin(annotation) is list and get_args(annotation) == (int,): - return st.lists(st.integers()) + return st.lists(st.integers(), max_size=2) raise NotImplementedError( f"No value strategy for return annotation {annotation!r}; " "supported: int, list[int]" ) -_UNARY_INT_FNS: list[Callable[[int], int]] = [ +_UNARY_NUM_FNS: list[Callable[[int], int]] = [ lambda x: x, lambda x: x + 1, lambda x: x - 1, @@ -30,7 +30,7 @@ def _value_strategy_for(annotation: Any) -> st.SearchStrategy[Any]: lambda x: 3 * x + 1, ] -_BINARY_INT_FNS: list[Callable[[int, int], int]] = [ +_BINARY_NUM_FNS: list[Callable[[int, int], int]] = [ lambda x, y: x + y, lambda x, y: x - y, lambda x, y: x * y, @@ -58,10 +58,10 @@ def _strategy_for_op(op: Operation) -> st.SearchStrategy[Callable[..., Any]]: if not params: return _value_strategy_for(ret).map(deffn) - if ret is int and param_types == (int,): - return st.sampled_from(_UNARY_INT_FNS) - if ret is int and param_types == (int, int): - return st.sampled_from(_BINARY_INT_FNS) + if ret in (int, float) and param_types == (int,): + return st.sampled_from(_UNARY_NUM_FNS) + if ret in (int, float) and param_types == (int, int): + return st.sampled_from(_BINARY_NUM_FNS) if get_origin(ret) is list and get_args(ret) == (int,) and param_types == (int,): return st.sampled_from(_UNARY_LIST_FNS) raise NotImplementedError( diff --git a/tests/test_handlers_llm_provider.py b/tests/test_handlers_llm_provider.py index bd2b6ccd2..0f5dd5f70 100644 --- a/tests/test_handlers_llm_provider.py +++ b/tests/test_handlers_llm_provider.py @@ -250,7 +250,7 @@ def test_agent_tool_names_are_valid_integration(): agent = _ToolNameAgent() template = agent.ask tools = template.tools - expected_helper_tool_name = f"self__{agent.helper.__name__}" + expected_helper_tool_name = "self__helper" assert tools assert expected_helper_tool_name in tools assert all(re.fullmatch(r"[a-zA-Z0-9_-]+", name) for name in tools) diff --git a/tests/test_ops_monoid.py b/tests/test_ops_monoid.py index a22928cca..e73a9a7b2 100644 --- a/tests/test_ops_monoid.py +++ b/tests/test_ops_monoid.py @@ -1,12 +1,23 @@ import functools import itertools +import typing import pytest from hypothesis import given, settings from hypothesis import strategies as st from effectful.internals.runtime import interpreter -from effectful.ops.monoid import Max, Min, NormalizeIntp, Product, Semilattice, Sum +from effectful.ops.monoid import ( + CartesianProduct, + Max, + Min, + Monoid, + NormalizeIntp, + Product, + Semilattice, + Sum, + distributes_over, +) from effectful.ops.semantics import apply, evaluate, fvsof, handler from effectful.ops.syntax import _BaseTerm, defdata, syntactic_eq from effectful.ops.types import NotHandled, Operation @@ -37,6 +48,18 @@ pytest.param(Product, id="Product"), ] +# Pairs (outer, inner) such that inner distributes over outer — i.e. the lifting +# identity ``outer(inner(body, A), CartesianProduct...) == inner(outer(body, D), ...)`` +# is valid for that semiring pair. +MONOID_PAIRS = [ + pytest.param(o.values[0], i.values[0], id=f"{o.id}-{i.id}") + for o in ALL_MONOIDS + for i in ALL_MONOIDS + if distributes_over( + typing.cast(Monoid, i.values[0]).plus, typing.cast(Monoid, o.values[0]).plus + ) +] + def define_vars(*names, typ=int): if len(names) == 1: @@ -70,14 +93,11 @@ def syntactic_eq_alpha(x, y) -> bool: def _canonicalize(expr): counter = itertools.count() - def _passthrough(op, *args, **kwargs): - return defdata(op, *args, **kwargs) - def _substitute(arg, renaming): """Apply a bound-variable renaming using ``evaluate`` for traversal.""" if not renaming: return arg - with interpreter({apply: _passthrough, **renaming}): + with interpreter({apply: _BaseTerm, **renaming}): return evaluate(arg) def _bound_var_order(args, kwargs, bound_set): @@ -121,7 +141,7 @@ def _apply_canonical(op, *args, **kwargs): *bindings.args, *bindings.kwargs.values() ) if not all_bound: - return defdata(op, *args, **kwargs) + return _BaseTerm(op, *args, **kwargs) order = _bound_var_order(args, kwargs, all_bound) canonical = {var: _canonical_op(next(counter)) for var in order} @@ -252,8 +272,8 @@ def test_plus_assoc_left(monoid): def test_plus_sequence(monoid): a, b, c, d = define_vars("a", "b", "c", "d", typ=type(monoid.identity)) _check_pair( - lhs=monoid.plus([a(), b()], [c(), d()]), - rhs=[monoid.plus(a(), c()), monoid.plus(b(), d())], + lhs=monoid.plus((a(), b()), (c(), d())), + rhs=(monoid.plus(a(), c()), monoid.plus(b(), d())), free_vars=[a, b, c, d], ) @@ -368,8 +388,8 @@ def f(_x: int) -> int: g = Operation.define(f, name="g") - lhs = monoid.reduce([f(x()), g(x())], {x: X()}) - rhs = [monoid.reduce(f(x()), {x: X()}), monoid.reduce(g(x()), {x: X()})] + lhs = monoid.reduce((f(x()), g(x())), {x: X()}) + rhs = (monoid.reduce(f(x()), {x: X()}), monoid.reduce(g(x()), {x: X()})) _check_pair(lhs=lhs, rhs=rhs, free_vars=[X, f, g]) @@ -385,11 +405,11 @@ def f(_x: int) -> int: g = Operation.define(f, name="g") - lhs = monoid.reduce([f(x()), g(y())], {x: X(), y: Y()}) - rhs = [ + lhs = monoid.reduce((f(x()), g(y())), {x: X(), y: Y()}) + rhs = ( monoid.reduce(f(x()), {x: X(), y: Y()}), monoid.reduce(g(y()), {x: X(), y: Y()}), - ] + ) _check_pair(lhs=lhs, rhs=rhs, free_vars=[X, Y, f, g]) @@ -496,8 +516,6 @@ def g(_x: int) -> list[int]: Sum.reduce(Product.plus(b(), f(b(), c())), {b: g(a()), c: C()}), ) assert fvsof(bogus_rhs) != fvsof(lhs) - # Structural-only negative check: the normalizer correctly refused to apply - # the bogus factorization. assert not syntactic_eq_alpha(lhs, bogus_rhs) @@ -516,3 +534,101 @@ def f(_x: int, _y: int) -> int: Sum.reduce(Product.plus(b(), f(b(), c())), {b: B(), c: C()}), ) _check_pair(lhs=lhs, rhs=rhs, free_vars=[A, B, C, f]) + + +@pytest.mark.parametrize("outer,inner", MONOID_PAIRS) +def test_reduce_lifted_1(outer, inner): + a, i = define_vars("a", "i") + A, N, A_domain = define_vars("A", "N", "A_domain", typ=list[int]) + + @Operation.define + def f(_: int) -> float: + raise NotHandled + + term1 = outer.reduce( + inner.reduce(f(a()), {a: A()}), + {A: CartesianProduct.reduce(A_domain(), {i: N()})}, + ) + term2 = inner.reduce(outer.reduce(f(a()), {a: A_domain()}), {i: N()}) + _check_pair(lhs=term1, rhs=term2, free_vars=[N, A_domain, f]) + + +def test_reduce_cartesian_1(): + a, i = define_vars("a", "i") + A = define_vars("A", typ=list[int]) + + term1 = Sum.reduce( + Product.reduce(a(), {a: []}), + {A: CartesianProduct.reduce([], {i: []})}, + ) + term2 = Product.reduce(Sum.reduce(a(), {a: []}), {i: []}) + assert term1 == term2 + + +def test_reduce_cartesian_2(): + a, i = define_vars("a", "i") + A = define_vars("A", typ=list[int]) + + term1 = Sum.reduce( + Product.reduce(a(), {a: A()}), + {A: CartesianProduct.reduce([(0,)], {i: [0]})}, + ) + term2 = Product.reduce(Sum.reduce(a(), {a: [0]}), {i: [0]}) + assert term1 == term2 + + +@pytest.mark.parametrize("outer,inner", MONOID_PAIRS) +def test_reduce_lifted_multi_index(outer, inner): + a, i, j = define_vars("a", "i", "j") + A, N, M, A_domain = define_vars("A", "N", "M", "A_domain", typ=list[int]) + + @Operation.define + def f(_: int) -> float: + raise NotHandled + + term1 = outer.reduce( + inner.reduce(f(a()), {a: A()}), + {A: CartesianProduct.reduce(A_domain(), {i: N(), j: M()})}, + ) + term2 = inner.reduce( + outer.reduce(f(a()), {a: A_domain()}), + {i: N(), j: M()}, + ) + _check_pair(lhs=term1, rhs=term2, free_vars=[N, M, A_domain, f]) + + +@pytest.mark.parametrize("outer,inner", MONOID_PAIRS) +def test_reduce_lifted_2(outer, inner): + """The worked example on page 396 of 'Lifted Variable Elimination: + Decoupling the Operators from the Constraint Language'. + + """ + a, i, s, t = define_vars("a", "i", "s", "t") + A, N, T = define_vars("A", "N", "T", typ=list[int]) + + @Operation.define + def A_domain(_i: int) -> list[int]: + raise NotHandled + + @Operation.define + def f1(_a: int, _s: int) -> float: + raise NotHandled + + @Operation.define + def f2(_t: int, _a: int) -> float: + raise NotHandled + + term1 = outer.reduce( + inner.reduce(inner.plus(f1(a(), s()), f2(t(), a())), {a: A()}), + {A: CartesianProduct.reduce(A_domain(i()), {i: N()}), t: T()}, + ) + + term2 = outer.reduce( + inner.reduce( + outer.reduce(inner.plus(f1(a(), s()), f2(t(), a())), {a: A_domain(i())}), + {i: N()}, + ), + {t: T()}, + ) + + _check_pair(lhs=term1, rhs=term2, free_vars=[a, i, s, t, A, N, T, A_domain, f1, f2]) diff --git a/tests/test_ops_syntax.py b/tests/test_ops_syntax.py index a5fdb749c..ccfd2dae0 100644 --- a/tests/test_ops_syntax.py +++ b/tests/test_ops_syntax.py @@ -489,7 +489,6 @@ def _(self, x: bool) -> bool: ) assert isinstance(term_float, Term) - assert term_float.op.__name__ == "my_singledispatch" assert term_float.args == (1.5,) assert term_float.kwargs == {} From 794a6044606374f7fbbfb011f1ef8b56588bff9f Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Tue, 12 May 2026 16:42:40 -0400 Subject: [PATCH 03/16] Refactor `monoid.py` to remove class structure (#661) * Add monoid module (#653) * add monoid module * clean up * fix doctest * fix * wip * remove incorrect rule * add disjoint set tests and fix bug * lint * drop jax monoid defs * drop incorrect comment * add assert * reduce nondeterminism and add assertions * fix inconsistent stream numbering and missing constant factors * wip * cleanup * fix rule * wip * fix bug * cleanup * lin * wip * fix tests * format * lint * wip --- effectful/ops/monoid.py | 255 ++++++++++++++++++--------------------- effectful/ops/types.py | 76 ++++++++++++ tests/test_ops_monoid.py | 6 +- 3 files changed, 197 insertions(+), 140 deletions(-) diff --git a/effectful/ops/monoid.py b/effectful/ops/monoid.py index ad83de47b..0d6e230c0 100644 --- a/effectful/ops/monoid.py +++ b/effectful/ops/monoid.py @@ -10,20 +10,25 @@ from typing import Annotated, Any from effectful.internals.disjoint_set import DisjointSet -from effectful.internals.runtime import interpreter from effectful.ops.semantics import coproduct, evaluate, fvsof, fwd, handler from effectful.ops.syntax import ( ObjectInterpretation, Scoped, _NumberTerm, - defdata, deffn, implements, iter_, syntactic_eq, syntactic_hash, ) -from effectful.ops.types import Expr, Interpretation, NotHandled, Operation, Term +from effectful.ops.types import ( + Expr, + Interpretation, + NotHandled, + Operation, + Term, + _CustomSingleDispatchMethod, +) # Note: The streams value type should be something like Iterable[T], but some of # our target stream types (e.g. jax.Array) are not subtypes of Iterable @@ -64,60 +69,57 @@ def __init__(self, kernel: Callable[[T, T], T], identity: T): def __repr__(self): return f"{type(self)}({self.kernel}, {self.identity})" + def __eq__(self, other): + return id(self) == id(other) + + def __hash__(self): + return hash(id(self)) + @Operation.define - def plus[S: Body[T]](self, *args: S) -> S: + @_CustomSingleDispatchMethod + def plus[S](self, dispatch, *args: S) -> S: """Monoid addition with broadcasting over common collection types, callables, and interpretations. - """ if not args: return typing.cast(S, self.identity) + return dispatch(type(args[0]))(self, *args) + @plus.register(object) # type: ignore[attr-defined] + def _(self, *args): if any(isinstance(x, Term) for x in args): - return typing.cast(S, defdata(self.plus, *args)) + raise NotHandled + return functools.reduce(self.kernel, args, self.identity) - return self._plus(*args) - - @functools.singledispatchmethod - def _plus[S](self, *args: S) -> S: - return typing.cast(S, functools.reduce(self.kernel, args, self.identity)) - - @_plus.register(tuple) + @plus.register(tuple) # type: ignore[attr-defined] def _(self, *args): return tuple(self.plus(*vs) for vs in zip(*args, strict=True)) - @_plus.register(Generator) + @plus.register(Generator) # type: ignore[attr-defined] def _(self, *args): return (self.plus(*vs) for vs in zip(*args, strict=True)) - @_plus.register(Mapping) + @plus.register(Mapping) # type: ignore[attr-defined] def _(self, *args): if isinstance(args[0], Interpretation): keys = args[0].keys() - for b in args[1:]: if not isinstance(b, Interpretation): raise TypeError(f"Expected interpretation but got {b}") - - b_keys = b.keys() - if not keys == b_keys: + if not keys == b.keys(): raise ValueError( - f"Expected interpretation of {keys} but got {b_keys}" + f"Expected interpretation of {keys} but got {b.keys()}" ) - - result = {k: self.plus(*(handler(b)(b[k]) for b in args)) for k in keys} - return result + return {k: self.plus(*(handler(b)(b[k]) for b in args)) for k in keys} for b in args[1:]: if not isinstance(b, Mapping): raise TypeError(f"Expected mapping but got {b}") - all_values = collections.defaultdict(list) for d in args: for k, v in d.items(): all_values[k].append(v) - result = {k: self.plus(*vs) for (k, vs) in all_values.items()} - return result + return {k: self.plus(*vs) for (k, vs) in all_values.items()} @Operation.define @functools.singledispatchmethod @@ -155,12 +157,9 @@ def generator(loop_order) -> Iterator[Interpretation]: yield coproduct(intp, intp2) loop_order = list(order_streams(streams)) - try: - return self.plus( - *(handler(intp)(evaluate)(body) for intp in generator(loop_order)) - ) - except NotHandled: - return typing.cast(U, defdata(self.reduce, body, streams)) + return self.plus( + *(handler(intp)(evaluate)(body) for intp in generator(loop_order)) + ) @reduce.register # type: ignore[attr-defined] def _(self, body: Mapping, streams): @@ -175,35 +174,19 @@ def _(self, body: Generator, streams): return (self.reduce(x, streams) for x in body) -class IdempotentMonoid[T](Monoid[T]): - @Operation.define - def plus[S: Body[T]](self, *args: S) -> S: - return super().plus(*args) - - @Operation.define - def reduce[A, B, U: Body]( - self, - body: Annotated[U, Scoped[A | B]], - streams: Annotated[Streams, Scoped[A]], - ) -> Annotated[U, Scoped[B]]: - return super().reduce(body, streams) +def _is_monoid_plus(op: Operation) -> bool: + """True if ``op`` is the ``plus`` operation of some :class:`Monoid`.""" + owner = getattr(op, "__self__", None) + return isinstance(owner, Monoid) and op is owner.plus -class CommutativeMonoid[T](Monoid[T]): - @Operation.define - def plus[S: Body[T]](self, *args: S) -> S: - return super().plus(*args) - - @Operation.define - def reduce[A, B, U: Body]( - self, - body: Annotated[U, Scoped[A | B]], - streams: Annotated[Streams, Scoped[A]], - ) -> Annotated[U, Scoped[B]]: - return super().reduce(body, streams) +def _is_monoid_reduce(op: Operation) -> bool: + """True if ``op`` is the ``reduce`` operation of some :class:`Monoid`.""" + owner = getattr(op, "__self__", None) + return isinstance(owner, Monoid) and op is owner.reduce -class CommutativeMonoidWithZero[T](CommutativeMonoid[T]): +class MonoidWithZero[T](Monoid[T]): zero: T def __init__(self, kernel: Callable[[T, T], T], identity: T, zero: T): @@ -213,32 +196,6 @@ def __init__(self, kernel: Callable[[T, T], T], identity: T, zero: T): def __repr__(self): return f"{type(self)}({self.kernel}, {self.identity}, {self.zero})" - @Operation.define - def plus[S: Body[T]](self, *args: S) -> S: - return super().plus(*args) - - @Operation.define - def reduce[A, B, U: Body]( - self, - body: Annotated[U, Scoped[A | B]], - streams: Annotated[Streams, Scoped[A]], - ) -> Annotated[U, Scoped[B]]: - return super().reduce(body, streams) - - -class Semilattice[T](IdempotentMonoid[T], CommutativeMonoid[T]): - @Operation.define - def plus[S: Body[T]](self, *args: S) -> S: - return super().plus(*args) - - @Operation.define - def reduce[A, B, U: Body]( - self, - body: Annotated[U, Scoped[A | B]], - streams: Annotated[Streams, Scoped[A]], - ) -> Annotated[U, Scoped[B]]: - return super().reduce(body, streams) - @Operation.define def _arg_min[T]( @@ -271,15 +228,30 @@ def to_tuple(x): return [to_tuple(x) + to_tuple(y) for (x, y) in itertools.product(a, b)] -Min = Semilattice(kernel=min, identity=float("inf")) -Max = Semilattice(kernel=max, identity=float("-inf")) +Min = Monoid(kernel=min, identity=float("inf")) +Max = Monoid(kernel=max, identity=float("-inf")) ArgMin = Monoid(kernel=_arg_min, identity=(float("inf"), None)) ArgMax = Monoid(kernel=_arg_max, identity=(float("-inf"), None)) -Sum = CommutativeMonoid(kernel=_NumberTerm.__add__, identity=0) -Product = CommutativeMonoidWithZero(kernel=_NumberTerm.__mul__, identity=1, zero=0) +Sum = Monoid(kernel=_NumberTerm.__add__, identity=0) +Product = MonoidWithZero(kernel=_NumberTerm.__mul__, identity=1, zero=0) CartesianProduct = Monoid(kernel=product, identity=[()]) +@dataclass +class _ExtensiblePredicate[T]: + elems: set[T] + + def register(self, t: T) -> None: + self.elems.add(t) + + def __call__(self, t: T) -> bool: + return t in self.elems + + +is_commutative = _ExtensiblePredicate({Max, Min, Sum, Product}) +is_idempotent = _ExtensiblePredicate({Max, Min}) + + @dataclass class _ExtensibleBinaryRelation[S, T]: tuples: set[tuple[S, T]] @@ -292,13 +264,7 @@ def __call__(self, s: S, t: T) -> bool: distributes_over = _ExtensibleBinaryRelation( - { - (Max.plus, Min.plus), - (Min.plus, Max.plus), - (Sum.plus, Min.plus), - (Sum.plus, Max.plus), - (Product.plus, Sum.plus), - } + {(Max, Min), (Min, Max), (Sum, Min), (Sum, Max), (Product, Sum)} ) @@ -337,10 +303,12 @@ class PlusAssoc(ObjectInterpretation): @implements(Monoid.plus) def plus(self, monoid, *args): - if any(isinstance(x, Term) and x.op is monoid.plus for x in args): + def is_nested_plus(x): + return isinstance(x, Term) and x.op is monoid.plus + + if any(is_nested_plus(x) for x in args): flat_args = itertools.chain.from_iterable( - t.args if isinstance(t, Term) and t.op is monoid.plus else (t,) - for t in args + t.args if is_nested_plus(t) else (t,) for t in args ) assert len(args) > 0 return monoid.plus(*flat_args) @@ -353,33 +321,36 @@ class PlusDistr(ObjectInterpretation): @implements(Monoid.plus) def plus(self, monoid, *args): if any( - isinstance(x, Term) and distributes_over(monoid.plus, x.op) for x in args + isinstance(x, Term) + and _is_monoid_plus(x.op) + and distributes_over(monoid, x.op.__self__) + for x in args ): non_terms = [] - # group terms by head operation - by_head_op = defaultdict(list) + # group terms by their monoid + by_monoid: dict[Monoid, list[Term]] = defaultdict(list) for t in args: - if isinstance(t, Term): - by_head_op[t.op].append(t) + if isinstance(t, Term) and _is_monoid_plus(t.op): + by_monoid[t.op.__self__].append(t) else: non_terms.append(t) # distribute over each group progress = False final_sum = [] - for op, terms in by_head_op.items(): + for m, terms in by_monoid.items(): if ( len(terms) > 1 - and distributes_over(monoid.plus, op) - and not distributes_over(op, monoid.plus) + and distributes_over(monoid, m) + and not distributes_over(m, monoid) ): progress = True term_args = (t.args for t in terms) dist_terms = ( monoid.plus(*args) for args in itertools.product(*term_args) ) - final_sum.append(op(*dist_terms)) + final_sum.append(m.plus(*dist_terms)) else: final_sum += terms if progress: @@ -390,8 +361,10 @@ def plus(self, monoid, *args): class PlusZero(ObjectInterpretation): """x₁ * ... * 0 * ... * xₙ = 0""" - @implements(CommutativeMonoidWithZero.plus) + @implements(Monoid.plus) def plus(self, monoid, *args): + if not (isinstance(monoid, MonoidWithZero)): + return fwd() if any(x is monoid.zero for x in args): return monoid.zero return fwd() @@ -400,8 +373,11 @@ def plus(self, monoid, *args): class PlusConsecutiveDups(ObjectInterpretation): """x ⊕ x ⊕ y = x ⊕ y""" - @implements(IdempotentMonoid.plus) + @implements(Monoid.plus) def plus(self, monoid, *args): + if not is_idempotent(monoid): + return fwd() + dedup_args = ( args[i] for i in range(len(args)) @@ -423,8 +399,11 @@ def __eq__(self, other): def __hash__(self): return syntactic_hash(self) - @implements(Semilattice.plus) + @implements(Monoid.plus) def plus(self, monoid, *args): + if not (is_idempotent(monoid) and is_commutative(monoid)): + return fwd() + # elim dups args_count = Counter(self._HashableTerm(t) for t in args) if len(args_count) < len(args): @@ -475,7 +454,7 @@ class ReduceFusion(ObjectInterpretation): @implements(Monoid.reduce) def reduce(self, monoid, body, streams): - if isinstance(body, Term) and body.op == monoid.reduce: + if isinstance(body, Term) and body.op is monoid.reduce: return monoid.reduce(body.args[0], streams | body.args[1]) return fwd() @@ -485,9 +464,11 @@ class ReduceSplit(ObjectInterpretation): reduce(R, S, b1 + ... + bn) = reduce(R, S, b1) + ... + reduce(R, S, bn) """ - @implements(CommutativeMonoid.reduce) + @implements(Monoid.reduce) def reduce(self, monoid, body, streams): - if isinstance(body, Term) and body.op == monoid.plus: + if not is_commutative(monoid): + return fwd() + if isinstance(body, Term) and body.op is monoid.plus: return monoid.plus(*(monoid.reduce(x, streams) for x in body.args)) return fwd() @@ -506,9 +487,16 @@ class ReduceFactorization(ObjectInterpretation): and free(Aᵢ) ∩ S ⊆ Sᵢ """ - @implements(CommutativeMonoid.reduce) + @implements(Monoid.reduce) def reduce(self, monoid, body, streams): - if isinstance(body, Term) and distributes_over(body.op, monoid.plus): + if not is_commutative(monoid): + return fwd() + if ( + isinstance(body, Term) + and _is_monoid_plus(body.op) + and distributes_over(body.op.__self__, monoid) + ): + inner_monoid: Monoid = body.op.__self__ stream_vars = set(streams.keys()) factors = [(arg, fvsof(arg)) for arg in body.args] stream_ids = {v: i for (i, v) in enumerate(stream_vars)} @@ -521,7 +509,7 @@ def reduce(self, monoid, body, streams): ds.union(stream_id, *deps) # factors are in the same partition as their dependencies - for factor, factor_fvs in factors: + for _, factor_fvs in factors: factor_streams = sorted( [stream_ids[v] for v in (factor_fvs & stream_vars)] ) @@ -550,14 +538,14 @@ def reduce(self, monoid, body, streams): for t in partition_factors ), "partition contains all streams required by factor" - partition_term = body.op(*(t[0] for t in partition_factors)) + partition_term = inner_monoid.plus(*(t[0] for t in partition_factors)) new_reduces.append((partition_term, partition_streams)) placed_streams |= partition_stream_keys constant_factors = [t for (t, fvs) in factors if not (fvs & stream_vars)] if len(new_reduces) > 1: - result = body.op( + result = inner_monoid.plus( *constant_factors, *(monoid.reduce(*args) for args in new_reduces) ) return result @@ -592,18 +580,6 @@ def inner_stream( ) -def match_reduce(term: Term) -> tuple | None: - reduce_args = None - - def set_reduce_args(*args, **kwargs): - nonlocal reduce_args - reduce_args = args - - with interpreter({Monoid.reduce: set_reduce_args}): - term.op(*term.args, **term.kwargs) - return reduce_args - - class ReduceDistributeCartesianProduct(ObjectInterpretation): """Eliminates a reduce over a cartesian product. ∑_x₁ ∑_x₂ ... ∑_xₙ ∏_i f(xᵢ) = ∏_i ∑_xᵢ f(xᵢ) @@ -623,25 +599,30 @@ class ReduceDistributeCartesianProduct(ObjectInterpretation): variable elimination." AISTATS. 2013. """ - @implements(CommutativeMonoid.reduce) + @implements(Monoid.reduce) def reduce(self, sum_monoid: Monoid, sum_body, sum_streams): - if not (isinstance(sum_body, Term)): + if not (is_commutative(sum_monoid) and isinstance(sum_body, Term)): return fwd() # body is a product or multiplication of products - if distributes_over(sum_body.op, sum_monoid.plus): + if _is_monoid_plus(sum_body.op) and distributes_over( + sum_body.op.__self__, sum_monoid + ): prod_reduces = sum_body.args else: prod_reduces = [sum_body] products: list[tuple[Monoid, Callable, Operation, Term]] = [] for prod_reduce in prod_reduces: - prod_args = match_reduce(prod_reduce) - if prod_args is None: + if not ( + isinstance(prod_reduce, Term) and _is_monoid_reduce(prod_reduce.op) + ): return fwd() - (prod_monoid, prod_body, prod_streams) = prod_args + prod_monoid: Monoid = prod_reduce.op.__self__ + prod_body = prod_reduce.args[0] + prod_streams = typing.cast(Mapping, prod_reduce.args[1]) if not ( - distributes_over(prod_monoid.plus, sum_monoid.plus) + distributes_over(prod_monoid, sum_monoid) and (len(products) == 0 or products[-1][0] == prod_monoid) ): return fwd() @@ -658,7 +639,7 @@ def reduce(self, sum_monoid: Monoid, sum_body, sum_streams): for outer_sum_streams, cprod_op, cprod_term in inner_stream(sum_streams): if not ( isinstance(cprod_term, Term) - and cprod_term.op == CartesianProduct.reduce + and cprod_term.op is CartesianProduct.reduce ): continue (cprod_body, cprod_streams) = cprod_term.args diff --git a/effectful/ops/types.py b/effectful/ops/types.py index 34379554e..bd87d1abf 100644 --- a/effectful/ops/types.py +++ b/effectful/ops/types.py @@ -42,6 +42,59 @@ def __call__(self, *args: P.args, **kwargs: P.kwargs) -> T: return self.func(self.dispatch, *args, **kwargs) +class _CustomSingleDispatchMethod[**P, **Q, S, T]: + """Method analog of :class:`_CustomSingleDispatchCallable`. + + The wrapped function has signature ``(self, dispatch, *args, **kwargs)``, + where ``dispatch`` is :meth:`functools.singledispatch.dispatch`. As a + descriptor, it binds ``self`` on attribute access, so callers invoke it + as ``instance.method(*args, **kwargs)``. + """ + + def __init__( + self, + func: Callable[Concatenate[Any, Callable[[type], Callable[Q, S]], P], T], + ): + self.func = func + self._registry = functools.singledispatch(func) + self.__signature__ = inspect.signature( + functools.partial(func, None, None) # type: ignore[arg-type] + ) + functools.update_wrapper(self, func) # type: ignore[arg-type] + + @property + def dispatch(self): + return self._registry.dispatch + + @property + def register(self): + return self._registry.register + + def __get__(self, instance, owner=None): + if instance is None: + return self + return _BoundCustomSingleDispatchMethod(self, instance) + + +class _BoundCustomSingleDispatchMethod: + __slots__ = ("_method", "_instance") + + def __init__(self, method: _CustomSingleDispatchMethod, instance: Any): + self._method = method + self._instance = instance + + @property + def dispatch(self): + return self._method.dispatch + + @property + def register(self): + return self._method.register + + def __call__(self, *args, **kwargs): + return self._method.func(self._instance, self._method.dispatch, *args, **kwargs) + + class _ClassMethodOpDescriptor(classmethod): def __init__(self, define, *args, **kwargs): super().__init__(*args, **kwargs) @@ -319,6 +372,15 @@ def func(*args, **kwargs): return typing.cast(Operation[P, T], cls.define(func, **kwargs)) + @define.register(types.MethodType) + @classmethod + def _define_methodtype[**P, T]( + cls, t: Callable[P, T], *, name: str | None = None + ) -> "Operation[P, T]": + op = cls._define_callable(t, name=name) + op.__self__ = t.__self__ # type: ignore[attr-defined] + return typing.cast("Operation[P, T]", op) + @define.register(staticmethod) @classmethod def _define_staticmethod[**P, T](cls, t: "staticmethod[P, T]", **kwargs): @@ -358,6 +420,20 @@ def func(*args, **kwargs): op.register = default._registry.register # type: ignore[attr-defined] return op + @define.register(_CustomSingleDispatchMethod) + @classmethod + def _define_customsingledispatchmethod( + cls, default: _CustomSingleDispatchMethod, **kwargs + ): + @functools.wraps(default.func) + def _wrapper(obj, *args, **kwargs): + return default.__get__(obj)(*args, **kwargs) + + op = cls.define(_wrapper, **kwargs) + op.register = default.register # type: ignore[attr-defined] + op.dispatch = default.dispatch # type: ignore[attr-defined] + return op + @typing.final def __default_rule__(self, *args: Q.args, **kwargs: Q.kwargs) -> "Expr[V]": """The default rule is used when the operation is not handled. diff --git a/tests/test_ops_monoid.py b/tests/test_ops_monoid.py index e73a9a7b2..d881869ac 100644 --- a/tests/test_ops_monoid.py +++ b/tests/test_ops_monoid.py @@ -14,9 +14,9 @@ Monoid, NormalizeIntp, Product, - Semilattice, Sum, distributes_over, + is_commutative, ) from effectful.ops.semantics import apply, evaluate, fvsof, handler from effectful.ops.syntax import _BaseTerm, defdata, syntactic_eq @@ -56,7 +56,7 @@ for o in ALL_MONOIDS for i in ALL_MONOIDS if distributes_over( - typing.cast(Monoid, i.values[0]).plus, typing.cast(Monoid, o.values[0]).plus + typing.cast(Monoid, i.values[0]), typing.cast(Monoid, o.values[0]) ) ] @@ -354,7 +354,7 @@ def test_plus_idempotent_non_consecutive(monoid): PlusDups; plain IdempotentMonoid leaves it as-is (consecutive-only).""" a, b = define_vars("a", "b") lhs = monoid.plus(a(), b(), a()) - if isinstance(monoid, Semilattice): + if is_commutative(monoid): rhs = monoid.plus(a(), b()) else: rhs = monoid.plus(a(), b(), a()) From effec0be1039b3573ad21fdd34bd9c26441684e5 Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Thu, 21 May 2026 15:11:59 -0400 Subject: [PATCH 04/16] Add `jax` array monoids and reduction rule (#658) * Add monoid module (#653) * add monoid module * clean up * fix doctest * fix * wip * remove incorrect rule * add disjoint set tests and fix bug * lint * drop jax monoid defs * drop incorrect comment * add assert * reduce nondeterminism and add assertions * fix inconsistent stream numbering and missing constant factors * wip * cleanup * wip * fix rule * wip * fix bug * cleanup * lin * wip * fix tests * format * lint * wip * wip * wip * wip * wip * wip * wip * wip * drop runtime typed dict lifting * wip * format * reorganize * stop using string dicts to avoid unification issue * wip * wip * wip * wip * wip * use check_rewrite in jax tests * lint * fix bugs --- effectful/handlers/jax/_handlers.py | 7 + effectful/handlers/jax/monoid.py | 162 +++++++ effectful/ops/monoid.py | 486 +++++++++++-------- effectful/ops/syntax.py | 2 + tests/_monoid_helpers.py | 284 ++++++++++- tests/test_handlers_jax_monoid.py | 96 ++++ tests/test_ops_monoid.py | 718 +++++++++++++++------------- 7 files changed, 1206 insertions(+), 549 deletions(-) create mode 100644 effectful/handlers/jax/monoid.py create mode 100644 tests/test_handlers_jax_monoid.py diff --git a/effectful/handlers/jax/_handlers.py b/effectful/handlers/jax/_handlers.py index 9c933af43..6779c8f44 100644 --- a/effectful/handlers/jax/_handlers.py +++ b/effectful/handlers/jax/_handlers.py @@ -22,6 +22,7 @@ deffn, defop, syntactic_eq, + syntactic_hash, ) from effectful.ops.types import Expr, NotHandled, Operation, Term @@ -328,3 +329,9 @@ def _(x: jax.Array, other) -> bool: and x.shape == other.shape and bool((jnp.asarray(x) == jnp.asarray(other)).all()) ) + + +@syntactic_hash.register(jax.Array) +def _(x: jax.Array) -> int: + # Concrete arrays aren't hashable; hash by shape, dtype, and bytes. + return hash(("jax.Array", x.shape, str(x.dtype), bytes(jax.numpy.asarray(x)))) diff --git a/effectful/handlers/jax/monoid.py b/effectful/handlers/jax/monoid.py new file mode 100644 index 000000000..a406cda5b --- /dev/null +++ b/effectful/handlers/jax/monoid.py @@ -0,0 +1,162 @@ +import functools + +import jax + +import effectful.handlers.jax.numpy as jnp +from effectful.handlers.jax import bind_dims, unbind_dims +from effectful.handlers.jax.scipy.special import logsumexp +from effectful.ops.monoid import ( + CartesianProduct, + Max, + Min, + Monoid, + NormalizeIntp, + Product, + Sum, + outer_stream, +) +from effectful.ops.semantics import evaluate, fvsof, fwd, handler, typeof +from effectful.ops.syntax import ObjectInterpretation, deffn, implements +from effectful.ops.types import Operation + + +def cartesian_prod(x, y): + if x.ndim == 1: + x = x[:, None] + if y.ndim == 1: + y = y[:, None] + nx, dx = x.shape + ny, dy = y.shape + # Broadcast into (nx, ny, dx+dy), then flatten the first two axes + x_b = jnp.broadcast_to(x[:, None, :], (nx, ny, dx)) + y_b = jnp.broadcast_to(y[None, :, :], (nx, ny, dy)) + return jnp.concatenate([x_b, y_b], axis=-1).reshape(nx * ny, dx + dy) + + +LogSumExp = Monoid(name="LogSumExp", identity=jnp.asarray(float("-inf"))) + + +def _jax_args(args): + """True iff ``args`` is non-empty and every arg is a concrete + :class:`jax.Array` (no Terms). + """ + typs = (typeof(a) for a in args) + return ( + bool(args) + and any(issubclass(t, jax.Array) for t in typs) + and all(issubclass(t, jax.typing.ArrayLike) for t in typs) + ) + + +class SumPlusJax(ObjectInterpretation): + @implements(Sum.plus) + def plus(self, *args): + if not _jax_args(args): + return fwd() + return functools.reduce(jnp.add, args) + + +class ProductPlusJax(ObjectInterpretation): + @implements(Product.plus) + def plus(self, *args): + if not _jax_args(args): + return fwd() + return functools.reduce(jnp.multiply, args) + + +class MinPlusJax(ObjectInterpretation): + @implements(Min.plus) + def plus(self, *args): + if not _jax_args(args): + return fwd() + return functools.reduce(jnp.minimum, args) + + +class MaxPlusJax(ObjectInterpretation): + @implements(Max.plus) + def plus(self, *args): + if not _jax_args(args): + return fwd() + return functools.reduce(jnp.maximum, args) + + +class LogSumExpPlusJax(ObjectInterpretation): + @implements(LogSumExp.plus) + def plus(self, *args): + if not _jax_args(args): + return fwd() + return functools.reduce(jnp.logaddexp, args) + + +class CartesianProductPlusJax(ObjectInterpretation): + @implements(CartesianProduct.plus) + def plus(self, *args): + # Skip identity ``[()]`` args; short-circuit on zero ``[]``. Both + # sentinels arrive as Python lists alongside jax-array factors, so + # check for them explicitly before composing. + if not any(isinstance(a, jax.Array) for a in args): + return fwd() + result = None + for a in args: + if a is CartesianProduct.zero: + return CartesianProduct.zero + if a is CartesianProduct.identity: + continue + if not isinstance(a, jax.Array): + return fwd() + result = a if result is None else cartesian_prod(result, a) + return result if result is not None else CartesianProduct.identity + + +ARRAY_REDUCTORS = { + Sum: jnp.sum, + Product: jnp.prod, + Min: jnp.min, + Max: jnp.max, + LogSumExp: logsumexp, +} + + +class ArrayReduce(ObjectInterpretation): + @implements(Monoid.reduce) + def reduce(self, monoid, body, streams): + if monoid not in ARRAY_REDUCTORS or typeof(body) is not jax.Array: + return fwd() + if not streams: + return monoid.identity + + reductor = ARRAY_REDUCTORS[monoid] + index = Operation.define(jax.Array) + for stream_key, stream_body, streams_tail in outer_stream(streams): + if not issubclass(typeof(stream_body), jax.Array): + continue + + if stream_key in fvsof(body): + with handler({stream_key: deffn(unbind_dims(stream_body, index))}): + eval_body = evaluate(body) + eval_streams_tail = evaluate(streams_tail) + assert isinstance(eval_streams_tail, dict) + reduce_tail = ( + monoid.reduce(eval_body, eval_streams_tail) + if len(eval_streams_tail) > 0 + else eval_body + ) + return reductor(bind_dims(reduce_tail, index), axis=0) + else: + # TODO: In this case, the stream is unused in the body. The body + # should be multiplied by the length of the stream. The current + # behavior is not efficient. + return fwd() + + return fwd() + + +NormalizeIntp.extend( + ArrayReduce(), + SumPlusJax(), + ProductPlusJax(), + MinPlusJax(), + MaxPlusJax(), + LogSumExpPlusJax(), + CartesianProductPlusJax(), +) diff --git a/effectful/ops/monoid.py b/effectful/ops/monoid.py index 0d6e230c0..70bb50022 100644 --- a/effectful/ops/monoid.py +++ b/effectful/ops/monoid.py @@ -1,10 +1,10 @@ import collections.abc import functools import itertools -import numbers +import operator import typing -from collections import Counter, defaultdict -from collections.abc import Callable, Generator, Iterable, Iterator, Mapping +from collections import Counter, UserDict, defaultdict +from collections.abc import Callable, Generator, Iterable, Mapping from dataclasses import dataclass from graphlib import TopologicalSorter from typing import Annotated, Any @@ -14,21 +14,13 @@ from effectful.ops.syntax import ( ObjectInterpretation, Scoped, - _NumberTerm, deffn, implements, iter_, syntactic_eq, syntactic_hash, ) -from effectful.ops.types import ( - Expr, - Interpretation, - NotHandled, - Operation, - Term, - _CustomSingleDispatchMethod, -) +from effectful.ops.types import Expr, Interpretation, NotHandled, Operation, Term # Note: The streams value type should be something like Iterable[T], but some of # our target stream types (e.g. jax.Array) are not subtypes of Iterable @@ -43,31 +35,35 @@ ) -def order_streams[T](streams: Streams[T]) -> Iterable[tuple[Operation[[], T], Any]]: - """Determine an order to evaluate the streams based on their dependencies""" +def outer_stream( + streams: Streams, +) -> Iterable[tuple[Operation, Expr, dict[Operation, Expr]]]: + """Returns the streams that can be ordered outermost in the loop nest as + well as the remaining streams in the nest. + + """ stream_vars = set(streams.keys()) - dependencies = {k: fvsof(v) & stream_vars for k, v in streams.items()} - topo = TopologicalSorter(dependencies) + pred = {k: fvsof(v) & stream_vars for k, v in streams.items()} + topo = TopologicalSorter(pred) topo.prepare() - while topo.is_active(): - node_group = topo.get_ready() - for op in sorted(node_group): - yield (op, streams[op]) - topo.done(*node_group) + return ( + (op, streams[op], {k: v for (k, v) in streams.items() if k != op}) + for op in topo.get_ready() + ) class Monoid[T]: - kernel: Operation[[T, T], T] + """A monoid with ``plus`` and ``reduce`` :class:`Operation` s.""" + + _name: str identity: T - def __init__(self, kernel: Callable[[T, T], T], identity: T): + def __init__(self, identity: T, name: str): + self._name = name self.identity = identity - self.kernel = ( - kernel if isinstance(kernel, Operation) else Operation.define(kernel) - ) def __repr__(self): - return f"{type(self)}({self.kernel}, {self.identity})" + return f"Monoid({self._name!r})" def __eq__(self, other): return id(self) == id(other) @@ -75,166 +71,63 @@ def __eq__(self, other): def __hash__(self): return hash(id(self)) + # the weak typing allows us to write monoid.plus(monoid.identity, ) + # and monoid.plus(monoid.identity, ) @Operation.define - @_CustomSingleDispatchMethod - def plus[S](self, dispatch, *args: S) -> S: - """Monoid addition with broadcasting over common collection types, - callables, and interpretations. + def plus(self, *args: Any) -> Any: + """Monoid addition. Handlers supply per-monoid and broadcasting + behavior; the default rule only handles empty / Term cases. """ if not args: - return typing.cast(S, self.identity) - return dispatch(type(args[0]))(self, *args) - - @plus.register(object) # type: ignore[attr-defined] - def _(self, *args): - if any(isinstance(x, Term) for x in args): - raise NotHandled - return functools.reduce(self.kernel, args, self.identity) - - @plus.register(tuple) # type: ignore[attr-defined] - def _(self, *args): - return tuple(self.plus(*vs) for vs in zip(*args, strict=True)) - - @plus.register(Generator) # type: ignore[attr-defined] - def _(self, *args): - return (self.plus(*vs) for vs in zip(*args, strict=True)) - - @plus.register(Mapping) # type: ignore[attr-defined] - def _(self, *args): - if isinstance(args[0], Interpretation): - keys = args[0].keys() - for b in args[1:]: - if not isinstance(b, Interpretation): - raise TypeError(f"Expected interpretation but got {b}") - if not keys == b.keys(): - raise ValueError( - f"Expected interpretation of {keys} but got {b.keys()}" - ) - return {k: self.plus(*(handler(b)(b[k]) for b in args)) for k in keys} - - for b in args[1:]: - if not isinstance(b, Mapping): - raise TypeError(f"Expected mapping but got {b}") - all_values = collections.defaultdict(list) - for d in args: - for k, v in d.items(): - all_values[k].append(v) - return {k: self.plus(*vs) for (k, vs) in all_values.items()} + return self.identity + raise NotHandled @Operation.define - @functools.singledispatchmethod def reduce[A, B, U: Body]( self, body: Annotated[U, Scoped[A | B]], streams: Annotated[Streams, Scoped[A]], ) -> Annotated[U, Scoped[B]]: - if callable(body): - return typing.cast(U, lambda *a, **k: self.reduce(body(*a, **k), streams)) - - def generator(loop_order) -> Iterator[Interpretation]: - if len(loop_order) == 0: - return - - stream_key = loop_order[0][0] - stream_values = evaluate(streams[stream_key]) - stream_values_iter = iter(stream_values) # type: ignore[arg-type] - - # If we try to iterate and get a term instead of a real - # iterator, give up + """Reduce ``body`` over ``streams``. Handlers supply per-monoid and + broadcasting behavior; the default rule only handles the empty-stream + case. + """ + for stream_key, stream_body, streams_tail in outer_stream(streams): + if isinstance(stream_body, Term): + continue + stream_values_iter = iter(stream_body) if isinstance(stream_values_iter, Term) and stream_values_iter.op is iter_: - raise NotHandled - - if len(loop_order) == 1: - for val in stream_values_iter: - yield {stream_key: functools.partial(lambda v: v, val)} - else: - for val in stream_values_iter: - intp: Interpretation = { - stream_key: functools.partial(lambda v: v, val) - } - with handler(intp): - for intp2 in generator(loop_order[1:]): - yield coproduct(intp, intp2) - - loop_order = list(order_streams(streams)) - return self.plus( - *(handler(intp)(evaluate)(body) for intp in generator(loop_order)) - ) - - @reduce.register # type: ignore[attr-defined] - def _(self, body: Mapping, streams): - return {k: self.reduce(v, streams) for (k, v) in body.items()} - - @reduce.register # type: ignore[attr-defined] - def _(self, body: tuple, streams): - return tuple(self.reduce(x, streams) for x in body) - - @reduce.register # type: ignore[attr-defined] - def _(self, body: Generator, streams): - return (self.reduce(x, streams) for x in body) - - -def _is_monoid_plus(op: Operation) -> bool: - """True if ``op`` is the ``plus`` operation of some :class:`Monoid`.""" - owner = getattr(op, "__self__", None) - return isinstance(owner, Monoid) and op is owner.plus - - -def _is_monoid_reduce(op: Operation) -> bool: - """True if ``op`` is the ``reduce`` operation of some :class:`Monoid`.""" - owner = getattr(op, "__self__", None) - return isinstance(owner, Monoid) and op is owner.reduce + continue + new_reduces = [] + for stream_val in stream_values_iter: + with handler({stream_key: deffn(stream_val)}): + eval_args = evaluate((body, streams_tail)) + assert isinstance(eval_args, tuple) + new_reduces.append( + self.reduce(*eval_args) if streams_tail else eval_args[0] + ) + return self.plus(*new_reduces) + raise NotHandled class MonoidWithZero[T](Monoid[T]): zero: T - def __init__(self, kernel: Callable[[T, T], T], identity: T, zero: T): - super().__init__(kernel, identity) + def __init__(self, name: str, identity: T, zero: T): + super().__init__(name=name, identity=identity) self.zero = zero - def __repr__(self): - return f"{type(self)}({self.kernel}, {self.identity}, {self.zero})" - - -@Operation.define -def _arg_min[T]( - a: tuple[numbers.Number, T | None], b: tuple[numbers.Number, T | None] -) -> tuple[numbers.Number, T | None]: - if isinstance(a[0], Term) or isinstance(b[0], Term): - raise NotHandled - return b if b[0] < a[0] else a # type: ignore - - -@Operation.define -def _arg_max[T]( - a: tuple[numbers.Number, T | None], b: tuple[numbers.Number, T | None] -) -> tuple[numbers.Number, T | None]: - if isinstance(a[0], Term) or isinstance(b[0], Term): - raise NotHandled - return b if b[0] > a[0] else a # type: ignore - - -@Operation.define -def product[T]( - a: Iterable[tuple[T, ...] | T], b: Iterable[tuple[T, ...] | T] -) -> Iterable[tuple[T, ...]]: - if isinstance(a, Term) or isinstance(b, Term): - raise NotHandled - - def to_tuple(x): - return x if isinstance(x, tuple) else (x,) - return [to_tuple(x) + to_tuple(y) for (x, y) in itertools.product(a, b)] - - -Min = Monoid(kernel=min, identity=float("inf")) -Max = Monoid(kernel=max, identity=float("-inf")) -ArgMin = Monoid(kernel=_arg_min, identity=(float("inf"), None)) -ArgMax = Monoid(kernel=_arg_max, identity=(float("-inf"), None)) -Sum = Monoid(kernel=_NumberTerm.__add__, identity=0) -Product = MonoidWithZero(kernel=_NumberTerm.__mul__, identity=1, zero=0) -CartesianProduct = Monoid(kernel=product, identity=[()]) +Min = Monoid(name="Min", identity=float("inf")) +Max = Monoid(name="Max", identity=-float("inf")) +ArgMin = Monoid(name="ArgMin", identity=(Min.identity, None)) +ArgMax = Monoid(name="ArgMax", identity=(Max.identity, None)) +Sum = Monoid(name="Sum", identity=0) +Product = MonoidWithZero(name="Product", identity=1, zero=0) +# CartesianProduct values are "two-level indexable" (rows × positions). The +# identity ``[()]`` is one row of zero positions (composing with it preserves +# shape); the zero ``[]`` is no rows (absorbs under product). +CartesianProduct = MonoidWithZero(name="CartesianProduct", identity=[()], zero=[]) @dataclass @@ -268,6 +161,18 @@ def __call__(self, s: S, t: T) -> bool: ) +def _is_monoid_plus(op: Operation) -> bool: + """True if ``op`` is the ``plus`` operation of some :class:`Monoid`.""" + owner = getattr(op, "__self__", None) + return isinstance(owner, Monoid) and op is owner.plus + + +def _is_monoid_reduce(op: Operation) -> bool: + """True if ``op`` is the ``reduce`` operation of some :class:`Monoid`.""" + owner = getattr(op, "__self__", None) + return isinstance(owner, Monoid) and op is owner.reduce + + class PlusEmpty(ObjectInterpretation): """plus() = 0""" @@ -319,7 +224,7 @@ class PlusDistr(ObjectInterpretation): """x + (y * z) = x * y + x * z""" @implements(Monoid.plus) - def plus(self, monoid, *args): + def plus(self, monoid: Monoid, *args): if any( isinstance(x, Term) and _is_monoid_plus(x.op) @@ -417,24 +322,6 @@ def plus(self, monoid, *args): return fwd() -NormalizePlusIntp = functools.reduce( - coproduct, - typing.cast( - list[Interpretation], - [ - PlusEmpty(), - PlusSingle(), - PlusIdentity(), - PlusAssoc(), - PlusDistr(), - PlusZero(), - PlusConsecutiveDups(), - PlusDups(), - ], - ), -) - - class ReduceNoStreams(ObjectInterpretation): """Implements the identity reduce(R, ∅, body) = 0 @@ -668,18 +555,217 @@ def reduce(self, sum_monoid: Monoid, sum_body, sum_streams): return fwd() -NormalizeReduceIntp = functools.reduce( - coproduct, - typing.cast( - list[Interpretation], - [ - ReduceNoStreams(), - ReduceFusion(), - ReduceSplit(), - ReduceFactorization(), - ReduceDistributeCartesianProduct(), - ], - ), +class MonoidOverCallable(ObjectInterpretation): + """``monoid.reduce(f, streams) = lambda *a: monoid.reduce(f(*a), streams)``.""" + + @implements(Monoid.reduce) + def reduce(self, monoid, body, streams): + if isinstance(body, Term) or not isinstance(body, Callable): + return fwd() + return lambda *a, **k: monoid.reduce(body(*a, **k), streams) + + @implements(Monoid.plus) + def plus(self, monoid, *args): + if not args or any( + isinstance(arg, Term) or not isinstance(arg, Callable) for arg in args + ): + return fwd() + return lambda *a, **k: monoid.plus(*(arg(*a, **k) for arg in args)) + + +class MonoidOverMapping(ObjectInterpretation): + """``monoid.reduce({k: v_k}, streams) = {k: monoid.reduce(v_k, streams)}``.""" + + @implements(Monoid.reduce) + def reduce(self, monoid, body, streams): + if isinstance(body, Term) or not isinstance(body, Mapping): + return fwd() + return {k: monoid.reduce(v, streams) for (k, v) in body.items()} + + @implements(Monoid.plus) + def plus(self, monoid, *args): + if not args or not isinstance(args[0], Mapping): + return fwd() + + if isinstance(args[0], Interpretation): + keys = args[0].keys() + for b in args[1:]: + if not isinstance(b, Interpretation): + raise TypeError(f"Expected interpretation but got {b}") + if not keys == b.keys(): + raise ValueError( + f"Expected interpretation of {keys} but got {b.keys()}" + ) + return {k: monoid.plus(*(handler(b)(b[k]) for b in args)) for k in keys} + + for b in args[1:]: + if not isinstance(b, Mapping): + raise TypeError(f"Expected mapping but got {b}") + all_values = collections.defaultdict(list) + for d in args: + for k, v in d.items(): + all_values[k].append(v) + return {k: monoid.plus(*vs) for (k, vs) in all_values.items()} + + +def _scalar_args(args): + """True iff ``args`` is non-empty and every arg is a concrete int/float.""" + return ( + bool(args) + and not any(isinstance(x, Term) for x in args) + and all(isinstance(x, int | float) for x in args) + ) + + +class SumPlus(ObjectInterpretation): + """Scalar implementation of :data:`Sum`.""" + + @implements(Sum.plus) + def plus(self, *args): + if not _scalar_args(args): + return fwd() + return sum(args) + + +class MinPlus(ObjectInterpretation): + """Scalar implementation of :data:`Min`.""" + + @implements(Min.plus) + def plus(self, *args): + if not _scalar_args(args): + return fwd() + return min(args) + + +class MaxPlus(ObjectInterpretation): + """Scalar implementation of :data:`Max`.""" + + @implements(Max.plus) + def plus(self, *args): + if not _scalar_args(args): + return fwd() + return max(args) + + +class ProductPlus(ObjectInterpretation): + """Scalar implementation of :data:`Product`.""" + + @implements(Product.plus) + def plus(self, *args): + if not _scalar_args(args): + return fwd() + return functools.reduce(operator.mul, args) + + +class ArgMinPlus(ObjectInterpretation): + """Scalar score implementation of :data:`ArgMin`.""" + + @implements(ArgMin.plus) + def plus(self, *args): + if not args or not all(isinstance(a, tuple) for a in args): + return fwd() + if any(isinstance(a[0], Term) for a in args): + return fwd() + if not all(isinstance(a[0], int | float) for a in args): + return fwd() + return min(args, key=lambda a: a[0]) + + +class ArgMaxPlus(ObjectInterpretation): + """Scalar score implementation of :data:`ArgMax`.""" + + @implements(ArgMax.plus) + def plus(self, *args): + if not args or not all(isinstance(a, tuple) for a in args): + return fwd() + if any(isinstance(a[0], Term) for a in args): + return fwd() + if not all(isinstance(a[0], int | float) for a in args): + return fwd() + return max(args, key=lambda a: a[0]) + + +class CartesianProductPlus(ObjectInterpretation): + """Pure-Python implementation of :data:`CartesianProduct`.""" + + @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): + return fwd() + + def to_tuple(x): + return x if isinstance(x, tuple) else (x,) + + return [ + sum((to_tuple(v) for v in vals), ()) for vals in itertools.product(*args) + ] + + +is_scalar = _ExtensiblePredicate({Min, Max, Sum, Product}) + + +class MonoidOverSequence(ObjectInterpretation): + @implements(Monoid.plus) + def plus(self, monoid, *args): + if ( + not is_scalar(monoid) + or not args + or not isinstance(args[0], tuple | list | Generator) + ): + return fwd() + zipped = zip(*args, strict=True) + result = (monoid.plus(*vs) for vs in zipped) + if isinstance(args[0], tuple | list): + return type(args[0])(result) + return result + + @implements(Monoid.reduce) + def reduce(self, monoid, body, streams): + if not is_scalar(monoid) or not isinstance(body, tuple | list | Generator): + return fwd() + result = (monoid.reduce(x, streams) for x in body) + if isinstance(body, tuple | list): + return type(body)(result) + return result + + +class _ExtensibleInterpretation(UserDict, Interpretation): + def extend(self, *intps: Interpretation) -> typing.Self: + for intp in intps: + self.data = coproduct(self.data, intp) # type: ignore[assignment] + return self + + +NormalizeIntp = _ExtensibleInterpretation().extend( + MonoidOverSequence(), + MonoidOverMapping(), + MonoidOverCallable(), + ReduceNoStreams(), + ReduceFusion(), + ReduceSplit(), + ReduceFactorization(), + ReduceDistributeCartesianProduct(), + PlusEmpty(), + PlusSingle(), + PlusIdentity(), + PlusAssoc(), + PlusDistr(), + PlusZero(), + PlusConsecutiveDups(), + PlusDups(), + SumPlus(), + MinPlus(), + MaxPlus(), + ProductPlus(), + ArgMinPlus(), + ArgMaxPlus(), + CartesianProductPlus(), ) +"""``NormalizeIntp``applies pure-Term rewrites (associativity, distributivity, +identity elimination, fusion, factorization, etc.). -NormalizeIntp = coproduct(NormalizePlusIntp, NormalizeReduceIntp) +""" diff --git a/effectful/ops/syntax.py b/effectful/ops/syntax.py index a2fbcf9b2..958f2fbf6 100644 --- a/effectful/ops/syntax.py +++ b/effectful/ops/syntax.py @@ -900,6 +900,8 @@ def _(x: collections.abc.Sequence, other) -> bool: @syntactic_eq.register(object) @syntactic_eq.register(str | bytes) def _(x: object, other) -> bool: + if isinstance(other, Term): # Terms often override __eq__ + return False return x == other diff --git a/tests/_monoid_helpers.py b/tests/_monoid_helpers.py index 9b311b257..f15103e30 100644 --- a/tests/_monoid_helpers.py +++ b/tests/_monoid_helpers.py @@ -1,23 +1,60 @@ +import itertools from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass from typing import Any, get_args, get_origin +import jax +from hypothesis import given, settings from hypothesis import strategies as st -from effectful.ops.syntax import deffn -from effectful.ops.types import Operation +import effectful.handlers.jax.numpy as _jnp +from effectful.internals.runtime import interpreter +from effectful.ops.monoid import NormalizeIntp +from effectful.ops.semantics import apply, evaluate, handler +from effectful.ops.syntax import _BaseTerm, defdata, deffn, syntactic_eq +from effectful.ops.types import NotHandled, Operation, Term + +_JAX_ARRAY_SHAPE = (2,) + + +def _jax_array_value_strategy() -> st.SearchStrategy[jax.Array]: + return st.lists( + st.integers(min_value=-5, max_value=5), + min_size=_JAX_ARRAY_SHAPE[0], + max_size=_JAX_ARRAY_SHAPE[0], + ).map(lambda xs: jax.numpy.asarray(xs, dtype=jax.numpy.float32)) + + +# Unary jax fns map a scalar to a 1-D array (analogous to ``_UNARY_LIST_FNS`` +# for ints). Uses the effectful-wrapped jnp so named-dim broadcasting works. +_UNARY_JAX_FNS: list[Callable[[jax.Array], jax.Array]] = [ + lambda a: _jnp.stack([a, a + 1]), + lambda a: _jnp.stack([a, -a]), + lambda a: _jnp.stack([a, a + 1, 2 * a]), +] + +_BINARY_JAX_FNS: list[Callable[[jax.Array, jax.Array], jax.Array]] = [ + lambda a, b: a + b, + lambda a, b: a - b, + lambda a, b: a * b, +] def _value_strategy_for(annotation: Any) -> st.SearchStrategy[Any]: """Strategy for the value an *0-arg* Operation should return.""" if annotation is int: - return st.integers() + return st.integers(min_value=-100, max_value=100) if annotation is float: return st.floats(allow_nan=False) if get_origin(annotation) is list and get_args(annotation) == (int,): - return st.lists(st.integers(), max_size=2) + return st.lists(st.integers(min_value=-100, max_value=100), max_size=2) + if annotation is jax.Array: + return _jax_array_value_strategy() + if get_origin(annotation) is list and get_args(annotation) == (jax.Array,): + return st.lists(_jax_array_value_strategy(), max_size=2) raise NotImplementedError( f"No value strategy for return annotation {annotation!r}; " - "supported: int, list[int]" + "supported: int, list[int], jax.Array, list[jax.Array]" ) @@ -46,6 +83,13 @@ def _value_strategy_for(annotation: Any) -> st.SearchStrategy[Any]: lambda x: [0, x, x + 1], ] +_UNARY_JAX_LIST_FNS: list[Callable[[jax.Array], list[jax.Array]]] = [ + lambda _x: [], + lambda x: [x], + lambda x: [x, x + 1], + lambda x: [x, -x], +] + def _strategy_for_op(op: Operation) -> st.SearchStrategy[Callable[..., Any]]: """Pick a strategy producing a callable suitable for binding `op` in an @@ -64,8 +108,18 @@ def _strategy_for_op(op: Operation) -> st.SearchStrategy[Callable[..., Any]]: return st.sampled_from(_BINARY_NUM_FNS) if get_origin(ret) is list and get_args(ret) == (int,) and param_types == (int,): return st.sampled_from(_UNARY_LIST_FNS) + if ret is jax.Array and param_types == (jax.Array,): + return st.sampled_from(_UNARY_JAX_FNS) + if ret is jax.Array and param_types == (jax.Array, jax.Array): + return st.sampled_from(_BINARY_JAX_FNS) + if ( + get_origin(ret) is list + and get_args(ret) == (jax.Array,) + and param_types == (jax.Array,) + ): + return st.sampled_from(_UNARY_JAX_LIST_FNS) raise NotImplementedError( - f"Function-typed free var must return int or list[int]; got {ret!r} for {op}" + f"No callable strategy for free var with return {ret!r}, params {param_types!r}" ) @@ -82,4 +136,220 @@ def random_interpretation( return intp -__all__ = ["random_interpretation"] +def define_vars(*names, typ=int): + if len(names) == 1: + return Operation.define(typ, name=names[0]) + return tuple(Operation.define(typ, name=n) for n in names) + + +def syntactic_eq_alpha(x, y) -> bool: + """Alpha-equivalence-respecting variant of ``syntactic_eq``. + + Walks each expression bottom-up with :func:`evaluate` and renames + every bound variable to a deterministic canonical Operation. The + canonical names are assigned by a counter that increments in + ``evaluate``'s natural traversal order, so two alpha-equivalent + expressions canonicalize to syntactically identical results. + """ + + _op_cache: dict[int, Operation] = {} + + def _canonical_op(idx: int, op: Operation) -> Operation: + """Cached canonical Operation, keyed by encounter index. + + Cached so that two independent canonicalize runs return the same + Operation object for the same index — letting ``syntactic_eq`` + compare canonical forms by Operation identity. + """ + if idx in _op_cache: + return _op_cache[idx] + + op = Operation.define(op, name=f"__cv_{idx}") + _op_cache[idx] = op + return op + + cx = _canonicalize(x, _canonical_op) + cy = _canonicalize(y, _canonical_op) + return syntactic_eq(cx, cy) + + +def _canonicalize(expr, _canonical_op): + counter = itertools.count() + + def _substitute(arg, renaming): + """Apply a bound-variable renaming using ``evaluate`` for traversal.""" + if not renaming: + return arg + with interpreter({apply: _BaseTerm, **renaming}): + return evaluate(arg) + + def _bound_var_order(args, kwargs, bound_set: set[Operation]) -> list[Operation]: + """Return bound variables in deterministic encounter order.""" + seen: list[Operation] = [] + seen_set: set[Operation] = set() + + def _capture(op, *a, **kw): + if op in bound_set and op not in seen_set: + seen.append(op) + seen_set.add(op) + return defdata(op, *a, **kw) + + # ``evaluate`` walks Terms, lists, tuples, mappings, dataclasses, + # etc. for free; the apply handler captures bound vars used as + # ``x()`` anywhere in the body. + with interpreter({apply: _capture}): + evaluate((args, kwargs)) + + # Binders bypass the apply handler. Pick them up with a small structural + # walk that visits dict keys too. + def _walk_bare(obj): + if isinstance(obj, Operation): + if obj in bound_set and obj not in seen_set: + seen.append(obj) + seen_set.add(obj) + elif isinstance(obj, dict): + for k, v in obj.items(): + _walk_bare(k) + _walk_bare(v) + elif isinstance(obj, list | set | frozenset | tuple): + for v in obj: + _walk_bare(v) + + _walk_bare((args, kwargs)) + return seen + + def _apply_canonical(op, *args, **kwargs) -> Term: + bindings = op.__fvs_rule__(*args, **kwargs) + all_bound: set[Operation] = set().union( + *bindings.args, *bindings.kwargs.values() + ) + if not all_bound: + return _BaseTerm(op, *args, **kwargs) + + order = _bound_var_order(args, kwargs, all_bound) + canonical = {var: _canonical_op(next(counter), var) for var in order} + assert all_bound <= set(order) + + new_args = tuple( + _substitute( + arg, {v: canonical[v] for v in bindings.args[i] if v in canonical} + ) + for i, arg in enumerate(args) + ) + new_kwargs = { + k: _substitute( + v, + {var: canonical[var] for var in bindings.kwargs[k] if var in canonical}, + ) + for k, v in kwargs.items() + } + + # avoid the renaming from defdata + return _BaseTerm(op, *new_args, **new_kwargs) + + with interpreter({apply: _apply_canonical}): + return evaluate(expr) + + +@dataclass(frozen=True) +class Backend: + """A value-domain spec used to share monoid tests across int and jax.Array + backends. Provides the concrete value type, the hypothesis strategy for + drawing scalars in property tests, and an equality predicate that works + for that domain. + """ + + name: str + scalar_typ: Any + stream_typ: Any + scalar_strategy: st.SearchStrategy[Any] + eq: Callable[[Any, Any], bool] + + def fresh_op(self, name: str, n_args: int = 1, ret: str = "scalar") -> Operation: + """Build a fresh, unhandled Operation whose parameter and return + annotations are derived from this backend. + + ``ret`` is ``"scalar"`` for a scalar return or ``"stream"`` for a + stream-of-scalar return. The operation has ``n_args`` parameters, + each of type ``scalar_typ``. + """ + scalar = self.scalar_typ + out = self.stream_typ if ret == "stream" else scalar + params = ", ".join(f"_a{i}" for i in range(n_args)) + ns: dict[str, Any] = {"NotHandled": NotHandled} + exec(f"def _fn({params}):\n raise NotHandled\n", ns) + fn = ns["_fn"] + fn.__annotations__ = { + **{f"_a{i}": scalar for i in range(n_args)}, + "return": out, + } + return Operation.define(fn, name=name) + + +def _int_eq(a: Any, b: Any) -> bool: + return not isinstance(a, Term) and not isinstance(b, Term) and a == b + + +def _jax_eq(a: Any, b: Any) -> bool: + def _leaf_eq(x: Any, y: Any) -> bool: + return bool(jax.numpy.all(jax.numpy.isclose(x, y, equal_nan=True))) + + try: + leaves = jax.tree.leaves(jax.tree.map(_leaf_eq, a, b)) + except (ValueError, TypeError): + return False + return all(leaves) + + +def check_rewrite( + lhs, + rhs, + rule, + *, + backend: Backend, + free_vars=[], + max_examples: int = 25, + deadline=None, +) -> None: + with handler(rule): + norm = evaluate(lhs) + assert syntactic_eq_alpha(norm, rhs) + + @given(intp=random_interpretation(free_vars)) + @settings(max_examples=max_examples, deadline=deadline) + def _check_semantics(intp): + with handler(NormalizeIntp), handler(intp): + lhs_val = evaluate(lhs) + rhs_val = evaluate(rhs) + assert backend.eq(lhs_val, rhs_val) + + _check_semantics() + + +INT_BACKEND = Backend( + name="int", + scalar_typ=int, + stream_typ=list[int], + scalar_strategy=st.integers(min_value=-100, max_value=100), + eq=_int_eq, +) + + +JAX_BACKEND = Backend( + name="jax", + scalar_typ=jax.Array, + stream_typ=jax.Array, + scalar_strategy=_jax_array_value_strategy(), + eq=_jax_eq, +) + + +__all__ = [ + "Backend", + "INT_BACKEND", + "JAX_BACKEND", + "random_interpretation", + "define_vars", + "syntactic_eq_alpha", + "check_rewrite", +] diff --git a/tests/test_handlers_jax_monoid.py b/tests/test_handlers_jax_monoid.py new file mode 100644 index 000000000..35d041fe2 --- /dev/null +++ b/tests/test_handlers_jax_monoid.py @@ -0,0 +1,96 @@ +import jax +import pytest + +import effectful.handlers.jax.numpy as jnp +from effectful.handlers.jax import bind_dims, unbind_dims +from effectful.handlers.jax.monoid import ArrayReduce, LogSumExp +from effectful.handlers.jax.scipy.special import logsumexp +from effectful.ops.monoid import Max, Min, Product, Sum +from tests._monoid_helpers import JAX_BACKEND, Backend, check_rewrite, define_vars + +MONOIDS = [ + pytest.param(Sum, jnp.sum, id="Sum"), + pytest.param(Product, jnp.prod, id="Product"), + pytest.param(Min, jnp.min, id="Min"), + pytest.param(Max, jnp.max, id="Max"), + pytest.param(LogSumExp, logsumexp, id="LogSumExp"), +] + + +@pytest.fixture +def backend() -> Backend: + return JAX_BACKEND + + +@pytest.mark.parametrize("monoid,reductor", MONOIDS) +def test_reduce_array_1(monoid, reductor, backend: Backend): + (x, k) = define_vars("x", "k", typ=jax.Array) + X = define_vars("X", typ=backend.stream_typ) + + lhs = monoid.reduce(x(), {x: X()}) + rhs = reductor(bind_dims(unbind_dims(X(), k), k), axis=0) + + check_rewrite( + lhs=lhs, rhs=rhs, rule=ArrayReduce(), backend=backend, free_vars=[x, X, k] + ) + + +@pytest.mark.parametrize("monoid,reductor", MONOIDS) +def test_reduce_array_2(monoid, reductor, backend: Backend): + (x, y, k1, k2) = define_vars("x", "y", "k1", "k2", typ=backend.scalar_typ) + (X, Y) = define_vars("X", "Y", typ=backend.stream_typ) + f = backend.fresh_op("f", n_args=2, ret="scalar") + + lhs = monoid.reduce(f(x(), y()), {x: X(), y: Y()}) + rhs = reductor( + bind_dims( + reductor( + bind_dims(f(unbind_dims(X(), k1), unbind_dims(Y(), k2)), k2), + axis=0, + ), + k1, + ), + axis=0, + ) + + check_rewrite( + lhs=lhs, + rhs=rhs, + rule=ArrayReduce(), + backend=backend, + free_vars=[x, y, k1, k2, X, Y, f], + ) + + +@pytest.mark.parametrize("monoid,reductor", MONOIDS) +def test_reduce_array_3(monoid, reductor, backend: Backend): + """Stream `y` is `g(x())` — depends on the bound element of X. The reducer + must inline ``g`` along the same named dim used to unbind `x`.""" + (x, y, k1, k2) = define_vars("x", "y", "k1", "k2", typ=backend.scalar_typ) + X = define_vars("X", typ=backend.stream_typ) + + f = backend.fresh_op("f", n_args=2, ret="scalar") + g = backend.fresh_op("g", n_args=1, ret="stream") + + lhs = monoid.reduce(f(x(), y()), {x: X(), y: g(x())}) + rhs = reductor( + bind_dims( + reductor( + bind_dims( + f(unbind_dims(X(), k1), unbind_dims(g(unbind_dims(X(), k1)), k2)), + k2, + ), + axis=0, + ), + k1, + ), + axis=0, + ) + + check_rewrite( + lhs=lhs, + rhs=rhs, + rule=ArrayReduce(), + backend=backend, + free_vars=[x, y, k1, k2, X, f, g], + ) diff --git a/tests/test_ops_monoid.py b/tests/test_ops_monoid.py index d881869ac..c7ee7567c 100644 --- a/tests/test_ops_monoid.py +++ b/tests/test_ops_monoid.py @@ -1,29 +1,51 @@ -import functools -import itertools import typing import pytest -from hypothesis import given, settings +from hypothesis import HealthCheck, given, settings from hypothesis import strategies as st -from effectful.internals.runtime import interpreter +import effectful.handlers.jax.monoid # noqa: F401 from effectful.ops.monoid import ( CartesianProduct, Max, Min, Monoid, + MonoidOverMapping, + MonoidOverSequence, NormalizeIntp, + PlusAssoc, + PlusConsecutiveDups, + PlusDistr, + PlusDups, + PlusEmpty, + PlusIdentity, + PlusSingle, + PlusZero, Product, + ReduceDistributeCartesianProduct, + ReduceFactorization, + ReduceFusion, + ReduceNoStreams, + ReduceSplit, Sum, distributes_over, - is_commutative, ) -from effectful.ops.semantics import apply, evaluate, fvsof, handler -from effectful.ops.syntax import _BaseTerm, defdata, syntactic_eq -from effectful.ops.types import NotHandled, Operation -from tests._monoid_helpers import random_interpretation +from effectful.ops.semantics import fvsof, handler +from effectful.ops.types import Operation +from tests._monoid_helpers import ( + INT_BACKEND, + JAX_BACKEND, + Backend, + check_rewrite, + define_vars, + syntactic_eq_alpha, +) + + +@pytest.fixture(params=[INT_BACKEND, JAX_BACKEND], ids=["int", "jax"]) +def backend(request) -> Backend: + return request.param -_INT = st.integers(min_value=-100, max_value=100) ALL_MONOIDS = [ pytest.param(Sum, id="Sum"), @@ -61,247 +83,183 @@ ] -def define_vars(*names, typ=int): - if len(names) == 1: - return Operation.define(typ, name=names[0]) - return tuple(Operation.define(typ, name=n) for n in names) - - -@functools.cache -def _canonical_op(idx: int) -> Operation: - """Globally cached canonical Operation, keyed by encounter index. - - Cached so that two independent canonicalize runs return the same - Operation object for the same index — letting ``syntactic_eq`` - compare canonical forms by Operation identity. - """ - return Operation.define(int, name=f"__cv_{idx}") - - -def syntactic_eq_alpha(x, y) -> bool: - """Alpha-equivalence-respecting variant of ``syntactic_eq``. - - Walks each expression bottom-up with :func:`evaluate` and renames - every bound variable to a deterministic canonical Operation. The - canonical names are assigned by a counter that increments in - ``evaluate``'s natural traversal order, so two alpha-equivalent - expressions canonicalize to syntactically identical results. - """ - return syntactic_eq(_canonicalize(x), _canonicalize(y)) - - -def _canonicalize(expr): - counter = itertools.count() - - def _substitute(arg, renaming): - """Apply a bound-variable renaming using ``evaluate`` for traversal.""" - if not renaming: - return arg - with interpreter({apply: _BaseTerm, **renaming}): - return evaluate(arg) - - def _bound_var_order(args, kwargs, bound_set): - """Return bound variables in deterministic encounter order.""" - seen: list[Operation] = [] - seen_set: set[Operation] = set() - - def _capture(op, *a, **kw): - if op in bound_set and op not in seen_set: - seen.append(op) - seen_set.add(op) - return defdata(op, *a, **kw) - - # ``evaluate`` walks Terms, lists, tuples, mappings, dataclasses, - # etc. for free; the apply handler captures bound vars used as - # ``x()`` anywhere in the body. - with interpreter({apply: _capture}): - evaluate((args, kwargs)) - - # Binders bypass the apply handler. Pick them up with a small structural - # walk that visits dict keys too. - def _walk_bare(obj): - if isinstance(obj, Operation): - if obj in bound_set and obj not in seen_set: - seen.append(obj) - seen_set.add(obj) - elif isinstance(obj, dict): - for k, v in obj.items(): - _walk_bare(k) - _walk_bare(v) - elif isinstance(obj, list | set | frozenset | tuple): - for v in obj: - _walk_bare(v) - - _walk_bare((args, kwargs)) - return seen - - def _apply_canonical(op, *args, **kwargs): - bindings = op.__fvs_rule__(*args, **kwargs) - all_bound: set[Operation] = set().union( - *bindings.args, *bindings.kwargs.values() - ) - if not all_bound: - return _BaseTerm(op, *args, **kwargs) - - order = _bound_var_order(args, kwargs, all_bound) - canonical = {var: _canonical_op(next(counter)) for var in order} - assert all_bound <= set(order) - - new_args = tuple( - _substitute( - arg, {v: canonical[v] for v in bindings.args[i] if v in canonical} - ) - for i, arg in enumerate(args) - ) - new_kwargs = { - k: _substitute( - v, - {var: canonical[var] for var in bindings.kwargs[k] if var in canonical}, - ) - for k, v in kwargs.items() - } - - # avoid the renaming from defdata - return _BaseTerm(op, *new_args, **new_kwargs) - - with interpreter({apply: _apply_canonical}): - return evaluate(expr) - - @pytest.mark.parametrize("monoid", ALL_MONOIDS) -@given(a=_INT, b=_INT, c=_INT) -@settings(max_examples=50, deadline=None) -def test_associativity(monoid, a, b, c): - left = monoid.plus(monoid.plus(a, b), c) - right = monoid.plus(a, monoid.plus(b, c)) - assert left == right +@given(data=st.data()) +@settings( + max_examples=50, + deadline=None, + suppress_health_check=[HealthCheck.function_scoped_fixture], +) +def test_associativity(monoid, backend, data): + a = data.draw(backend.scalar_strategy) + b = data.draw(backend.scalar_strategy) + c = data.draw(backend.scalar_strategy) + with handler(NormalizeIntp): + left = monoid.plus(monoid.plus(a, b), c) + right = monoid.plus(a, monoid.plus(b, c)) + assert backend.eq(left, right) @pytest.mark.parametrize("monoid", ALL_MONOIDS) -@given(a=_INT) -@settings(max_examples=50, deadline=None) -def test_identity(monoid, a): - assert monoid.plus(monoid.identity, a) == a - assert monoid.plus(a, monoid.identity) == a +@given(data=st.data()) +@settings( + max_examples=50, + deadline=None, + suppress_health_check=[HealthCheck.function_scoped_fixture], +) +def test_identity(monoid, backend, data): + a = data.draw(backend.scalar_strategy) + with handler(NormalizeIntp): + assert backend.eq(monoid.plus(monoid.identity, a), a) + assert backend.eq(monoid.plus(a, monoid.identity), a) @pytest.mark.parametrize("monoid", COMMUTATIVE) -@given(a=_INT, b=_INT) -@settings(max_examples=50, deadline=None) -def test_commutativity(monoid, a, b): - assert monoid.plus(a, b) == monoid.plus(b, a) +@given(data=st.data()) +@settings( + max_examples=50, + deadline=None, + suppress_health_check=[HealthCheck.function_scoped_fixture], +) +def test_commutativity(monoid, backend, data): + a = data.draw(backend.scalar_strategy) + b = data.draw(backend.scalar_strategy) + with handler(NormalizeIntp): + assert backend.eq(monoid.plus(a, b), monoid.plus(b, a)) @pytest.mark.parametrize("monoid", IDEMPOTENT) -@given(a=_INT) -@settings(max_examples=50, deadline=None) -def test_idempotence(monoid, a): - assert monoid.plus(a, a) == a +@given(data=st.data()) +@settings( + max_examples=50, + deadline=None, + suppress_health_check=[HealthCheck.function_scoped_fixture], +) +def test_idempotence(monoid, backend, data): + a = data.draw(backend.scalar_strategy) + with handler(NormalizeIntp): + assert backend.eq(monoid.plus(a, a), a) @pytest.mark.parametrize("monoid", WITH_ZERO) -@given(a=_INT) -@settings(max_examples=50, deadline=None) -def test_zero_absorbs(monoid, a): - assert monoid.plus(monoid.zero, a) == monoid.zero - assert monoid.plus(a, monoid.zero) == monoid.zero - - -def _check_pair(lhs, rhs, *, free_vars=[], max_examples: int = 25) -> None: - """Run structural + semantic checks on a TermPair.""" +@given(data=st.data()) +@settings( + max_examples=50, + deadline=None, + suppress_health_check=[HealthCheck.function_scoped_fixture], +) +def test_zero_absorbs(monoid, backend, data): + a = data.draw(backend.scalar_strategy) with handler(NormalizeIntp): - norm = evaluate(lhs) - - assert syntactic_eq_alpha(norm, rhs) + assert backend.eq(monoid.plus(monoid.zero, a), monoid.zero) + assert backend.eq(monoid.plus(a, monoid.zero), monoid.zero) - @given(intp=random_interpretation(free_vars)) - @settings(max_examples=max_examples, deadline=None) - def _check_semantics(intp): - with handler(intp): - lhs_val = evaluate(lhs) - rhs_val = evaluate(rhs) - assert lhs_val == rhs_val - _check_semantics() +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_plus_empty(monoid, backend): + check_rewrite( + lhs=monoid.plus(), rhs=monoid.identity, rule=PlusEmpty(), backend=backend + ) @pytest.mark.parametrize("monoid", ALL_MONOIDS) -def test_plus_empty(monoid): - _check_pair(lhs=monoid.plus(), rhs=monoid.identity) +def test_plus_single(monoid, backend): + x = define_vars("x", typ=backend.scalar_typ) + check_rewrite( + lhs=monoid.plus(x()), rhs=x(), rule=PlusSingle(), backend=backend, free_vars=[x] + ) @pytest.mark.parametrize("monoid", ALL_MONOIDS) -def test_plus_single(monoid): - x = define_vars("x", typ=type(monoid.identity)) - _check_pair(lhs=monoid.plus(x()), rhs=x(), free_vars=[x]) +def test_plus_identity_right(monoid, backend): + x = define_vars("x", typ=backend.scalar_typ) + lhs = monoid.plus(x(), monoid.identity) + rhs = monoid.plus(x()) -@pytest.mark.parametrize("monoid", ALL_MONOIDS) -def test_plus_identity_right(monoid): - x = define_vars("x", typ=type(monoid.identity)) - _check_pair(lhs=monoid.plus(x(), monoid.identity), rhs=x(), free_vars=[x]) + check_rewrite(lhs=lhs, rhs=rhs, rule=PlusIdentity(), backend=backend, free_vars=[x]) @pytest.mark.parametrize("monoid", ALL_MONOIDS) -def test_plus_identity_left(monoid): - x = define_vars("x", typ=type(monoid.identity)) - _check_pair(lhs=monoid.plus(monoid.identity, x()), rhs=x(), free_vars=[x]) +def test_plus_identity_left(monoid, backend): + x = define_vars("x", typ=backend.scalar_typ) + + lhs = monoid.plus(monoid.identity, x()) + rhs = monoid.plus(x()) + + check_rewrite(lhs=lhs, rhs=rhs, rule=PlusIdentity(), backend=backend, free_vars=[x]) @pytest.mark.parametrize("monoid", ALL_MONOIDS) -def test_plus_assoc_right(monoid): - x, y, z = define_vars("x", "y", "z", typ=type(monoid.identity)) - _check_pair( +def test_plus_assoc_right(monoid, backend): + x, y, z = define_vars("x", "y", "z", typ=backend.scalar_typ) + check_rewrite( lhs=monoid.plus(x(), monoid.plus(y(), z())), rhs=monoid.plus(x(), y(), z()), + rule=PlusAssoc(), + backend=backend, free_vars=[x, y, z], ) @pytest.mark.parametrize("monoid", ALL_MONOIDS) -def test_plus_assoc_left(monoid): - x, y, z = define_vars("x", "y", "z", typ=type(monoid.identity)) - _check_pair( +def test_plus_assoc_left(monoid, backend): + x, y, z = define_vars("x", "y", "z", typ=backend.scalar_typ) + check_rewrite( lhs=monoid.plus(monoid.plus(x(), y()), z()), rhs=monoid.plus(x(), y(), z()), + rule=PlusAssoc(), + backend=backend, free_vars=[x, y, z], ) @pytest.mark.parametrize("monoid", ALL_MONOIDS) -def test_plus_sequence(monoid): - a, b, c, d = define_vars("a", "b", "c", "d", typ=type(monoid.identity)) - _check_pair( +def test_plus_sequence(monoid, backend): + a, b, c, d = define_vars("a", "b", "c", "d", typ=backend.scalar_typ) + check_rewrite( lhs=monoid.plus((a(), b()), (c(), d())), rhs=(monoid.plus(a(), c()), monoid.plus(b(), d())), + rule=MonoidOverSequence(), + backend=backend, free_vars=[a, b, c, d], ) @pytest.mark.parametrize("monoid", ALL_MONOIDS) -def test_plus_mapping(monoid): - a, b, c, d = define_vars("a", "b", "c", "d", typ=type(monoid.identity)) - _check_pair( - lhs=monoid.plus({"x": a(), "y": b()}, {"x": c(), "z": d()}), - rhs={"x": monoid.plus(a(), c()), "y": b(), "z": d()}, +def test_plus_mapping(monoid, backend): + a, b, c, d = define_vars("a", "b", "c", "d", typ=backend.scalar_typ) + + lhs = monoid.plus({0: a(), 1: b()}, {0: c(), 2: d()}) + rhs = {0: monoid.plus(a(), c()), 1: monoid.plus(b()), 2: monoid.plus(d())} + + check_rewrite( + lhs=lhs, + rhs=rhs, + rule=MonoidOverMapping(), + backend=backend, free_vars=[a, b, c, d], ) -def test_plus_distributes(): - a, b, c, d = define_vars("a", "b", "c", "d") +def test_plus_distributes(backend): + a, b, c, d = define_vars("a", "b", "c", "d", typ=backend.scalar_typ) lhs = Product.plus(Sum.plus(a(), b()), Sum.plus(c(), d())) - rhs = Sum.plus( - Product.plus(a(), c()), - Product.plus(a(), d()), - Product.plus(b(), c()), - Product.plus(b(), d()), + rhs = Product.plus( + Sum.plus( + Product.plus(a(), c()), + Product.plus(a(), d()), + Product.plus(b(), c()), + Product.plus(b(), d()), + ) + ) + check_rewrite( + lhs=lhs, rhs=rhs, rule=PlusDistr(), backend=backend, free_vars=[a, b, c, d] ) - _check_pair(lhs=lhs, rhs=rhs, free_vars=[a, b, c, d]) -def test_plus_distributes_constant(): - a, b, c, d = define_vars("a", "b", "c", "d") +def test_plus_distributes_constant(backend): + a, b, c, d = define_vars("a", "b", "c", "d", typ=backend.scalar_typ) lhs = Product.plus(Sum.plus(a(), b()), Sum.plus(c(), d()), 5) rhs = Product.plus( 5, @@ -312,11 +270,13 @@ def test_plus_distributes_constant(): Product.plus(b(), d()), ), ) - _check_pair(lhs=lhs, rhs=rhs, free_vars=[a, b, c, d]) + check_rewrite( + lhs=lhs, rhs=rhs, rule=PlusDistr(), backend=backend, free_vars=[a, b, c, d] + ) -def test_plus_distributes_multiple(): - a, b, c, d = define_vars("a", "b", "c", "d") +def test_plus_distributes_multiple(backend): + a, b, c, d = define_vars("a", "b", "c", "d", typ=backend.scalar_typ) lhs = Sum.plus( Min.plus(a(), b()), Min.plus(c(), d()), @@ -337,72 +297,123 @@ def test_plus_distributes_multiple(): Sum.plus(b(), d()), ), ) - _check_pair(lhs=lhs, rhs=rhs, free_vars=[a, b, c, d]) + check_rewrite( + lhs=lhs, rhs=rhs, rule=PlusDistr(), backend=backend, free_vars=[a, b, c, d] + ) @pytest.mark.parametrize("monoid", IDEMPOTENT) -def test_plus_idempotent_consecutive(monoid): +def test_plus_idempotent_consecutive(monoid, backend): """``a, a, b → a, b`` — only consecutive duplicates collapse.""" - a, b = define_vars("a", "b") + a, b = define_vars("a", "b", typ=backend.scalar_typ) lhs = monoid.plus(a(), a(), b()) - return _check_pair(lhs=lhs, rhs=monoid.plus(a(), b()), free_vars=[a, b]) + return check_rewrite( + lhs=lhs, + rhs=monoid.plus(a(), b()), + rule=PlusConsecutiveDups(), + backend=backend, + free_vars=[a, b], + ) @pytest.mark.parametrize("monoid", IDEMPOTENT) -def test_plus_idempotent_non_consecutive(monoid): +def test_plus_idempotent_non_consecutive(monoid, backend): """``a, b, a`` — Semilattice (Min/Max) collapses via commutative - PlusDups; plain IdempotentMonoid leaves it as-is (consecutive-only).""" - a, b = define_vars("a", "b") + PlusDups.""" + a, b = define_vars("a", "b", typ=backend.scalar_typ) lhs = monoid.plus(a(), b(), a()) - if is_commutative(monoid): - rhs = monoid.plus(a(), b()) - else: - rhs = monoid.plus(a(), b(), a()) - _check_pair(lhs=lhs, rhs=rhs, free_vars=[a, b]) + rhs = monoid.plus(a(), b()) + check_rewrite(lhs=lhs, rhs=rhs, rule=PlusDups(), backend=backend, free_vars=[a, b]) -def test_plus_commutative_idempotent_long(): +@pytest.mark.parametrize("monoid", [Min, Max]) +def test_plus_commutative_idempotent_long(monoid, backend): """Long alternation collapses via commutative dedup (Min/Max only).""" - a, b = define_vars("a", "b") - lhs = Min.plus(a(), b(), a(), b(), b(), a(), a()) - _check_pair(lhs=lhs, rhs=Min.plus(a(), b()), free_vars=[a, b]) + a, b = define_vars("a", "b", typ=backend.scalar_typ) + lhs = monoid.plus(a(), b(), a(), b(), b(), a(), a()) + rhs = monoid.plus(a(), b()) + check_rewrite(lhs=lhs, rhs=rhs, rule=PlusDups(), backend=backend, free_vars=[a, b]) @pytest.mark.parametrize("monoid", WITH_ZERO) -def test_plus_zero(monoid): - a = define_vars("a") +def test_plus_zero(monoid, backend): + a = define_vars("a", typ=backend.scalar_typ) lhs_right = monoid.plus(a(), monoid.zero) lhs_left = monoid.plus(monoid.zero, a()) - _check_pair(lhs=lhs_right, rhs=monoid.zero, free_vars=[a]) - _check_pair(lhs=lhs_left, rhs=monoid.zero, free_vars=[a]) + rhs = monoid.zero + check_rewrite( + lhs=lhs_right, rhs=rhs, rule=PlusZero(), backend=backend, free_vars=[a] + ) + check_rewrite( + lhs=lhs_left, rhs=rhs, rule=PlusZero(), backend=backend, free_vars=[a] + ) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_partial_1(monoid, backend): + x, y = define_vars("x", "y", typ=backend.scalar_typ) + lhs = monoid.reduce(x(), {x: []}) + rhs = monoid.identity + check_rewrite(lhs=lhs, rhs=rhs, rule={}, backend=backend, free_vars=[x, y]) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_partial_2(monoid, backend): + x, y = define_vars("x", "y", typ=backend.scalar_typ) + Y = define_vars("Y", typ=backend.stream_typ) + + lhs = monoid.reduce(x(), {y: Y(), x: []}) + rhs = monoid.identity + + check_rewrite(lhs=lhs, rhs=rhs, rule={}, backend=backend, free_vars=[x, y, Y]) @pytest.mark.parametrize("monoid", ALL_MONOIDS) -def test_reduce_body_sequence(monoid): - x = Operation.define(int, name="x") - X = Operation.define(list[int], name="X") +def test_partial_3(monoid, backend): + x, y, a, b = define_vars("x", "y", "a", "b", typ=backend.scalar_typ) + Y = define_vars("Y", typ=backend.stream_typ) + + lhs = monoid.reduce(x(), {y: Y(), x: [a(), b()]}) + rhs = monoid.plus(monoid.reduce(a(), {y: Y()}), monoid.reduce(b(), {y: Y()})) + + check_rewrite(lhs=lhs, rhs=rhs, rule={}, backend=backend, free_vars=[x, y, a, b, Y]) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_partial_4(monoid, backend): + x, y, a, b = define_vars("x", "y", "a", "b", typ=backend.scalar_typ) + f = backend.fresh_op("f", n_args=1, ret="stream") + + lhs = monoid.reduce(x(), {y: f(x()), x: [a(), b()]}) + rhs = monoid.plus(monoid.reduce(a(), {y: f(a())}), monoid.reduce(b(), {y: f(b())})) - @Operation.define - def f(_x: int) -> int: - raise NotHandled + check_rewrite(lhs=lhs, rhs=rhs, rule={}, backend=backend, free_vars=[x, y, a, b, f]) + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_reduce_body_sequence(monoid, backend): + x = Operation.define(backend.scalar_typ, name="x") + X = Operation.define(backend.stream_typ, name="X") + f = backend.fresh_op("f", n_args=1, ret="scalar") g = Operation.define(f, name="g") lhs = monoid.reduce((f(x()), g(x())), {x: X()}) rhs = (monoid.reduce(f(x()), {x: X()}), monoid.reduce(g(x()), {x: X()})) - _check_pair(lhs=lhs, rhs=rhs, free_vars=[X, f, g]) + check_rewrite( + lhs=lhs, + rhs=rhs, + rule=MonoidOverSequence(), + backend=backend, + free_vars=[X, f, g], + ) @pytest.mark.parametrize("monoid", ALL_MONOIDS) -def test_reduce_body_sequence_2(monoid): - x, y = define_vars("x", "y") - X, Y = define_vars("X", "Y", typ=list[int]) - - @Operation.define - def f(_x: int) -> int: - raise NotHandled - +def test_reduce_body_sequence_2(monoid, backend): + x, y = define_vars("x", "y", typ=backend.scalar_typ) + X, Y = define_vars("X", "Y", typ=backend.stream_typ) + f = backend.fresh_op("f", n_args=1, ret="scalar") g = Operation.define(f, name="g") lhs = monoid.reduce((f(x()), g(y())), {x: X(), y: Y()}) @@ -411,103 +422,115 @@ def f(_x: int) -> int: monoid.reduce(g(y()), {x: X(), y: Y()}), ) - _check_pair(lhs=lhs, rhs=rhs, free_vars=[X, Y, f, g]) + check_rewrite( + lhs=lhs, + rhs=rhs, + rule=MonoidOverSequence(), + backend=backend, + free_vars=[X, Y, f, g], + ) @pytest.mark.parametrize("monoid", ALL_MONOIDS) -def test_reduce_body_mapping(monoid): - x = Operation.define(int, name="x") - X = Operation.define(list[int], name="X") - - @Operation.define - def f(_x: int) -> int: - raise NotHandled - +def test_reduce_body_mapping(monoid, backend): + x = Operation.define(backend.scalar_typ, name="x") + X = Operation.define(backend.stream_typ, name="X") + f = backend.fresh_op("f", n_args=1, ret="scalar") g = Operation.define(f, name="g") - lhs = monoid.reduce({"a": f(x()), "b": g(x())}, {x: X()}) + lhs = monoid.reduce({0: f(x()), 1: g(x())}, {x: X()}) rhs = { - "a": monoid.reduce(f(x()), {x: X()}), - "b": monoid.reduce(g(x()), {x: X()}), + 0: monoid.reduce(f(x()), {x: X()}), + 1: monoid.reduce(g(x()), {x: X()}), } - _check_pair(lhs=lhs, rhs=rhs, free_vars=[X, f, g]) + check_rewrite( + lhs=lhs, + rhs=rhs, + rule=MonoidOverMapping(), + backend=backend, + free_vars=[X, f, g], + ) @pytest.mark.parametrize("monoid", ALL_MONOIDS) -def test_reduce_no_streams(monoid): - a = define_vars("a") +def test_reduce_no_streams(monoid, backend): + a = define_vars("a", typ=backend.scalar_typ) lhs = monoid.reduce(a(), {}) rhs = monoid.identity - _check_pair(lhs=lhs, rhs=rhs, free_vars=[a]) + check_rewrite( + lhs=lhs, rhs=rhs, rule=ReduceNoStreams(), backend=backend, free_vars=[a] + ) @pytest.mark.parametrize("monoid", ALL_MONOIDS) -def test_reduce_reduce(monoid): - a, b = define_vars("a", "b") - A, B = define_vars("A", "B", typ=list[int]) - - @Operation.define - def f(_x: int, _y: int) -> int: - raise NotHandled +def test_reduce_reduce(monoid, backend): + a, b = define_vars("a", "b", typ=backend.scalar_typ) + A, B = define_vars("A", "B", typ=backend.stream_typ) + f = backend.fresh_op("f", n_args=2, ret="scalar") lhs = monoid.reduce(monoid.reduce(f(a(), b()), {a: A()}), {b: B()}) rhs = monoid.reduce(f(a(), b()), {a: A(), b: B()}) - _check_pair(lhs=lhs, rhs=rhs, free_vars=[A, B, f]) + check_rewrite( + lhs=lhs, rhs=rhs, rule=ReduceFusion(), backend=backend, free_vars=[A, B, f] + ) @pytest.mark.parametrize("monoid", COMMUTATIVE) -def test_reduce_plus(monoid): - a, b = define_vars("a", "b") - A, B = define_vars("A", "B", typ=list[int]) +def test_reduce_plus(monoid, backend): + a, b = define_vars("a", "b", typ=backend.scalar_typ) + A, B = define_vars("A", "B", typ=backend.stream_typ) lhs = monoid.reduce(monoid.plus(a(), b()), {a: A(), b: B()}) rhs = monoid.plus( monoid.reduce(a(), {a: A(), b: B()}), monoid.reduce(b(), {a: A(), b: B()}), ) - _check_pair(lhs=lhs, rhs=rhs, free_vars=[A, B]) + check_rewrite( + lhs=lhs, rhs=rhs, rule=ReduceSplit(), backend=backend, free_vars=[A, B] + ) -def test_reduce_independent_1(): - a, b = define_vars("a", "b") - A, B = define_vars("A", "B", typ=list[int]) +def test_reduce_independent_1(backend): + a, b = define_vars("a", "b", typ=backend.scalar_typ) + A, B = define_vars("A", "B", typ=backend.stream_typ) lhs = Sum.reduce(Product.plus(a(), b()), {a: A(), b: B()}) - rhs = Product.plus(Sum.reduce(a(), {a: A()}), Sum.reduce(b(), {b: B()})) - _check_pair(lhs=lhs, rhs=rhs, free_vars=[A, B]) - + rhs = Product.plus( + Sum.reduce(Product.plus(a()), {a: A()}), Sum.reduce(Product.plus(b()), {b: B()}) + ) + check_rewrite( + lhs=lhs, rhs=rhs, rule=ReduceFactorization(), backend=backend, free_vars=[A, B] + ) -def test_reduce_independent_2(): - a, b, c = define_vars("a", "b", "c") - A, B, C = define_vars("A", "B", "C", typ=list[int]) - @Operation.define - def f(_x: int, _y: int) -> int: - raise NotHandled +def test_reduce_independent_2(backend): + a, b, c = define_vars("a", "b", "c", typ=backend.scalar_typ) + A, B, C = define_vars("A", "B", "C", typ=backend.stream_typ) + f = backend.fresh_op("f", n_args=2, ret="scalar") lhs = Sum.reduce(Product.plus(a(), b(), f(b(), c())), {a: A(), b: B(), c: C()}) rhs = Product.plus( - Sum.reduce(a(), {a: A()}), + Sum.reduce(Product.plus(a()), {a: A()}), Sum.reduce(Product.plus(b(), f(b(), c())), {b: B(), c: C()}), ) - _check_pair(lhs=lhs, rhs=rhs, free_vars=[A, B, C, f]) + check_rewrite( + lhs=lhs, + rhs=rhs, + rule=ReduceFactorization(), + backend=backend, + free_vars=[A, B, C, f], + ) -def test_reduce_independent_3_negative(): +def test_reduce_independent_3_negative(backend): """Stream `b` depends on `a` (b: g(a())), so the proposed factorization is unsound — the normalizer must NOT apply it.""" - a, b, c = define_vars("a", "b", "c") - A, C = define_vars("A", "C", typ=list[int]) - - @Operation.define - def f(_x: int, _y: int) -> int: - raise NotHandled + a, b, c = define_vars("a", "b", "c", typ=backend.scalar_typ) + A, C = define_vars("A", "C", typ=backend.stream_typ) + f = backend.fresh_op("f", n_args=2, ret="scalar") + g = backend.fresh_op("g", n_args=1, ret="stream") - @Operation.define - def g(_x: int) -> list[int]: - raise NotHandled - - with handler(NormalizeIntp): + with handler(ReduceFactorization()): # ty:ignore[invalid-argument-type] lhs = Sum.reduce( Product.plus(a(), b(), f(b(), c())), {a: A(), b: g(a()), c: C()} ) @@ -519,104 +542,107 @@ def g(_x: int) -> list[int]: assert not syntactic_eq_alpha(lhs, bogus_rhs) -def test_reduce_independent_4(): - a, b, c = define_vars("a", "b", "c") - A, B, C = define_vars("A", "B", "C", typ=list[int]) - - @Operation.define - def f(_x: int, _y: int) -> int: - raise NotHandled +def test_reduce_independent_4(backend): + a, b, c = define_vars("a", "b", "c", typ=backend.scalar_typ) + A, B, C = define_vars("A", "B", "C", typ=backend.stream_typ) + f = backend.fresh_op("f", n_args=2, ret="scalar") lhs = Sum.reduce(Product.plus(a(), b(), f(b(), c()), 7), {a: A(), b: B(), c: C()}) rhs = Product.plus( 7, - Sum.reduce(a(), {a: A()}), + Sum.reduce(Product.plus(a()), {a: A()}), Sum.reduce(Product.plus(b(), f(b(), c())), {b: B(), c: C()}), ) - _check_pair(lhs=lhs, rhs=rhs, free_vars=[A, B, C, f]) + check_rewrite( + lhs=lhs, + rhs=rhs, + rule=ReduceFactorization(), + backend=backend, + free_vars=[A, B, C, f], + ) @pytest.mark.parametrize("outer,inner", MONOID_PAIRS) -def test_reduce_lifted_1(outer, inner): - a, i = define_vars("a", "i") - A, N, A_domain = define_vars("A", "N", "A_domain", typ=list[int]) - - @Operation.define - def f(_: int) -> float: - raise NotHandled +def test_reduce_lifted_1(outer, inner, backend): + a, i = define_vars("a", "i", typ=backend.scalar_typ) + A, N, A_domain = define_vars("A", "N", "A_domain", typ=backend.stream_typ) + f = backend.fresh_op("f", n_args=1, ret="scalar") term1 = outer.reduce( inner.reduce(f(a()), {a: A()}), {A: CartesianProduct.reduce(A_domain(), {i: N()})}, ) - term2 = inner.reduce(outer.reduce(f(a()), {a: A_domain()}), {i: N()}) - _check_pair(lhs=term1, rhs=term2, free_vars=[N, A_domain, f]) + term2 = inner.reduce(outer.reduce(inner.plus(f(a())), {a: A_domain()}), {i: N()}) + + check_rewrite( + lhs=term1, + rhs=term2, + rule=ReduceDistributeCartesianProduct(), + backend=backend, + free_vars=[N, A_domain, f], + ) def test_reduce_cartesian_1(): - a, i = define_vars("a", "i") - A = define_vars("A", typ=list[int]) + a, i = define_vars("a", "i", typ=int) + A = define_vars("A", typ=tuple[int]) - term1 = Sum.reduce( - Product.reduce(a(), {a: []}), - {A: CartesianProduct.reduce([], {i: []})}, - ) - term2 = Product.reduce(Sum.reduce(a(), {a: []}), {i: []}) + with handler(NormalizeIntp): + term1 = Sum.reduce( + Product.reduce(a(), {a: []}), + {A: CartesianProduct.reduce([], {i: []})}, + ) + term2 = Product.reduce(Sum.reduce(a(), {a: []}), {i: []}) assert term1 == term2 def test_reduce_cartesian_2(): - a, i = define_vars("a", "i") - A = define_vars("A", typ=list[int]) + a, i = define_vars("a", "i", typ=int) + A = define_vars("A", typ=tuple[int]) - term1 = Sum.reduce( - Product.reduce(a(), {a: A()}), - {A: CartesianProduct.reduce([(0,)], {i: [0]})}, - ) - term2 = Product.reduce(Sum.reduce(a(), {a: [0]}), {i: [0]}) + with handler(NormalizeIntp): + term1 = Sum.reduce( + Product.reduce(a(), {a: A()}), + {A: CartesianProduct.reduce([(0,)], {i: [0]})}, + ) + term2 = Product.reduce(Sum.reduce(a(), {a: [0]}), {i: [0]}) assert term1 == term2 @pytest.mark.parametrize("outer,inner", MONOID_PAIRS) -def test_reduce_lifted_multi_index(outer, inner): - a, i, j = define_vars("a", "i", "j") - A, N, M, A_domain = define_vars("A", "N", "M", "A_domain", typ=list[int]) - - @Operation.define - def f(_: int) -> float: - raise NotHandled +def test_reduce_lifted_multi_index(outer, inner, backend): + a, i, j = define_vars("a", "i", "j", typ=backend.scalar_typ) + A, N, M, A_domain = define_vars("A", "N", "M", "A_domain", typ=backend.stream_typ) + f = backend.fresh_op("f", n_args=1, ret="scalar") term1 = outer.reduce( inner.reduce(f(a()), {a: A()}), {A: CartesianProduct.reduce(A_domain(), {i: N(), j: M()})}, ) term2 = inner.reduce( - outer.reduce(f(a()), {a: A_domain()}), + outer.reduce(inner.plus(f(a())), {a: A_domain()}), {i: N(), j: M()}, ) - _check_pair(lhs=term1, rhs=term2, free_vars=[N, M, A_domain, f]) + check_rewrite( + lhs=term1, + rhs=term2, + rule=ReduceDistributeCartesianProduct(), + backend=backend, + free_vars=[N, M, A_domain, f], + ) @pytest.mark.parametrize("outer,inner", MONOID_PAIRS) -def test_reduce_lifted_2(outer, inner): +def test_reduce_lifted_2(outer, inner, backend): """The worked example on page 396 of 'Lifted Variable Elimination: Decoupling the Operators from the Constraint Language'. """ - a, i, s, t = define_vars("a", "i", "s", "t") - A, N, T = define_vars("A", "N", "T", typ=list[int]) - - @Operation.define - def A_domain(_i: int) -> list[int]: - raise NotHandled - - @Operation.define - def f1(_a: int, _s: int) -> float: - raise NotHandled - - @Operation.define - def f2(_t: int, _a: int) -> float: - raise NotHandled + a, i, s, t = define_vars("a", "i", "s", "t", typ=backend.scalar_typ) + A, N, T = define_vars("A", "N", "T", typ=backend.stream_typ) + A_domain = backend.fresh_op("A_domain", n_args=1, ret="stream") + f1 = backend.fresh_op("f1", n_args=2, ret="scalar") + f2 = backend.fresh_op("f2", n_args=2, ret="scalar") term1 = outer.reduce( inner.reduce(inner.plus(f1(a(), s()), f2(t(), a())), {a: A()}), @@ -625,10 +651,18 @@ def f2(_t: int, _a: int) -> float: term2 = outer.reduce( inner.reduce( - outer.reduce(inner.plus(f1(a(), s()), f2(t(), a())), {a: A_domain(i())}), + outer.reduce( + inner.plus(inner.plus(f1(a(), s()), f2(t(), a()))), {a: A_domain(i())} + ), {i: N()}, ), {t: T()}, ) - _check_pair(lhs=term1, rhs=term2, free_vars=[a, i, s, t, A, N, T, A_domain, f1, f2]) + check_rewrite( + lhs=term1, + rhs=term2, + rule=ReduceDistributeCartesianProduct(), + backend=backend, + free_vars=[a, i, s, t, A, N, T, A_domain, f1, f2], + ) From e08ad5fe1f21bcc6e6c17f5bb794ad0f86f861ff Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Fri, 22 May 2026 12:51:04 -0400 Subject: [PATCH 05/16] Add `delta` terms for array construction in `handlers.jax.monoid` (#663) * Add monoid module (#653) * add monoid module * clean up * fix doctest * fix * wip * remove incorrect rule * add disjoint set tests and fix bug * lint * drop jax monoid defs * drop incorrect comment * add assert * reduce nondeterminism and add assertions * fix inconsistent stream numbering and missing constant factors * wip * cleanup * wip * wip * fix rule * wip * fix bug * cleanup * lin * wip * fix tests * format * lint * wip * wip * wip * wip * wip * wip * wip * wip * drop runtime typed dict lifting * wip * format * reorganize * stop using string dicts to avoid unification issue * wip * wip * wip * wip * wip * use check_rewrite in jax tests * lint * wip * fix bugs * comment on not implemented cases * format * simplify * lint * add matmul test --- effectful/handlers/jax/_handlers.py | 6 + effectful/handlers/jax/_terms.py | 20 +-- effectful/handlers/jax/monoid.py | 255 +++++++++++++++++++++++++++- effectful/ops/monoid.py | 6 +- tests/test_handlers_jax_monoid.py | 177 ++++++++++++++++++- 5 files changed, 442 insertions(+), 22 deletions(-) diff --git a/effectful/handlers/jax/_handlers.py b/effectful/handlers/jax/_handlers.py index 6779c8f44..920aa876e 100644 --- a/effectful/handlers/jax/_handlers.py +++ b/effectful/handlers/jax/_handlers.py @@ -138,6 +138,12 @@ def _partial_eval(t: Expr[jax.Array]) -> Expr[jax.Array]: if not sized_fvs: return t + # if any dimension is zero sized, the result is empty + if any(size == 0 for size in sized_fvs.values()): + key = tuple(sized_fvs.keys()) + shape = tuple(sized_fvs[k] for k in key) + return jax_getitem(jnp.empty(shape), key) + def _is_eager(t): return not isinstance(t, Term) or t.op in sized_fvs or is_eager_array(t) diff --git a/effectful/handlers/jax/_terms.py b/effectful/handlers/jax/_terms.py index 812062931..c88fe9341 100644 --- a/effectful/handlers/jax/_terms.py +++ b/effectful/handlers/jax/_terms.py @@ -8,7 +8,6 @@ import effectful.handlers.jax.numpy as jnp from effectful.handlers.jax._handlers import ( IndexElement, - _partial_eval, _register_jax_op, bind_dims, jax_getitem, @@ -451,28 +450,15 @@ def _bind_dims_array(t: jax.Array, *args: Operation[[], jax.Array]) -> jax.Array >>> bind_dims(t, b, a).shape (3, 2) """ - - def _evaluate(expr): - if isinstance(expr, Term): - (args, kwargs) = jax.tree.map(_evaluate, (expr.args, expr.kwargs)) - return _partial_eval(expr) - if not jax.tree_util.treedef_is_leaf(jax.tree.structure(expr)): - return jax.tree.map(_evaluate, expr) - return expr - if not isinstance(t, Term): return t - result = _evaluate(t) - if not isinstance(result, Term) or not args: - return result - # ensure that the result is a jax_getitem with an array as the first argument - if not (result.op is jax_getitem and isinstance(result.args[0], jax.Array)): + if not (t.op is jax_getitem and isinstance(t.args[0], jax.Array)): raise NotHandled - array = result.args[0] - dims = result.args[1] + array = t.args[0] + dims = t.args[1] assert isinstance(dims, Sequence) # ensure that the order is a subset of the named dimensions diff --git a/effectful/handlers/jax/monoid.py b/effectful/handlers/jax/monoid.py index a406cda5b..42d7866ec 100644 --- a/effectful/handlers/jax/monoid.py +++ b/effectful/handlers/jax/monoid.py @@ -1,4 +1,6 @@ import functools +import typing +from collections.abc import Iterable import jax @@ -12,12 +14,13 @@ Monoid, NormalizeIntp, Product, + Streams, Sum, outer_stream, ) from effectful.ops.semantics import evaluate, fvsof, fwd, handler, typeof from effectful.ops.syntax import ObjectInterpretation, deffn, implements -from effectful.ops.types import Operation +from effectful.ops.types import Interpretation, NotHandled, Operation, Term def cartesian_prod(x, y): @@ -151,8 +154,258 @@ def reduce(self, monoid, body, streams): return fwd() +@Operation.define +def delta(_index: tuple[int, ...], _weight: jax.Array) -> jax.Array: + raise NotHandled + + +py_range = range + + +@Operation.define +def range(*args: int) -> Iterable[jax.Array]: + raise NotHandled + + +def _range_start(term: Term): + assert term.op == range + if len(term.args) < 2: + return 0 + return term.args[0] + + +def _range_stop(term: Term): + assert term.op == range + if len(term.args) < 2: + return term.args[0] + return term.args[1] + + +def _range_step(term: Term): + assert term.op == range + if len(term.args) < 3: + return 1 + return term.args[2] + + +def _is_simple_range(term: Term) -> bool: + if term.op != range: + return False + + start = _range_start(term) + step = _range_step(term) + return ( + not isinstance(start, Term) + and start == 0 + and not isinstance(step, Term) + and step == 1 + ) + + +class ReduceDeltaIndependent(ObjectInterpretation): + """Eliminate a Delta that has independent, dense index arguments. + + reduce(M, streams, delta((), body)) ≡ reduce(M, streams, body) + + reduce(M, streams ∪ {v: range(N)}, delta(idx' ++ (v(),), body)) + ═══════════════════════════════════════════════════════════════════════════ + reduce(M, streams, delta(idx', bind_dims(body[v() := unbind_dims(streams[v], fv)], fv))) + + Not yet supported: + + - **Strided index streams** (``range(0, N, k)`` for ``k != 1``): the + premise ``_is_simple_range`` requires ``start == 0`` and ``step == 1``. + A strided extension would substitute ``v() := unbind_dims(jnp.arange( + start, stop, step), fv)`` and otherwise follow the same shape — the + change is purely in the recognised range form, the bind/unbind cycle + below is unchanged. + - **Non-zero start** (``range(a, b, 1)`` with ``a != 0``): same template + as the strided case; only the recognised range form changes. + - **Non-bare index expressions** (``delta((2*v(),), w)``, + ``delta((f(v()),), w)``, etc.): currently requires the final index + entry to be a bare call ``v()`` of a stream var op. Generalizing to + arbitrary index expressions is a scatter, not a bind: materialize the + index expression and the weight separately over ``v``, then + ``jnp.zeros(N).at[indices].set(values)`` (for Sum; analogous for + other monoids using ``.add``/``.min``/``.max``/...). This is a + different leaf operation from ``bind_dims`` and warrants a sibling + rule rather than an extension of this one. + """ + + @implements(Monoid.reduce) + def _(self, monoid: Monoid, body, streams: Streams): + if not (isinstance(body, Term) and body.op == delta): + return fwd() + + indices, weight = body.args + assert isinstance(indices, tuple) + + if not indices: + return monoid.reduce(weight, streams) + + head_indices, tail_index = indices[:-1], indices[-1] + if not (isinstance(tail_index, Term) and tail_index.op in streams): + return fwd() + + tail_op: Operation = tail_index.op + tail_stream = streams[tail_op] + if not (isinstance(tail_stream, Term) and _is_simple_range(tail_stream)): + return fwd() + + fresh_op = Operation.define(tail_op) + indices = jnp.arange(_range_stop(tail_stream)) + if isinstance(indices, jax.Array) and len(indices) == 0: + return monoid.identity + + fresh_stream = unbind_dims(indices, fresh_op) + subst_intp = typing.cast(Interpretation, {tail_op: deffn(fresh_stream)}) + fresh_body = bind_dims(handler(subst_intp)(evaluate)(weight), fresh_op) + fresh_streams = {k: v for (k, v) in streams.items() if k != tail_op} + return monoid.reduce(delta(head_indices, fresh_body), fresh_streams) + + +class ReduceDependentRangeMask(ObjectInterpretation): + """Eliminate a dependent range by masking. + + reduce(M, streams ∪ {u: range(N), v: range(u())}, body) + ═══════════════════════════════════════════════════════════════════════════ + reduce(M, streams ∪ {u: range(N), v: range(N)}, where(v() < u(), body, M.identity)) + + Currently recognises only the lower-triangular form ``v: range(u())``: + constant start of 0, dependent stop equal to a bare call of another + stream var. + + Not yet supported: + + - **Upper-triangular** (``v: range(u(), N)`` — constant stop, dependent + start): bbox becomes ``range(0, N)`` (or ``range(0, bbox_N)``), guard + becomes ``v() >= u()``. Same shape of rewrite as lower-tri; differs + only in which side of the range carries the stream-var reference and + in the predicate direction. + - **Banded** (``v: range(u() - k, u() + k + 1)`` — two-sided dependent + bounds with constant width): bbox is ``range(0, N + k)`` (or similar + bounded by both endpoints' extents), guard is + ``(v() >= u() - k) & (v() < u() + k + 1)``. Needs both-sides + affine-bound recognition. + - **Strided dependent** (``v: range(0, u(), k)`` for ``k != 1``): bbox + stays ``range(0, N)`` and guard becomes + ``(v() < u()) & (v() % k == 0)`` (or equivalent), or alternatively + embed in a smaller bbox ``range(0, ceil(N/k))`` and remap the index. + - **Affine bounds** (``v: range(a*u() + b, c*u() + d)`` for affine + coefficients): bbox computed from ``ub(c*u() + d)`` over ``u``'s + range; guard is the conjunction of the two affine constraints. This + subsumes the upper/banded/strided cases under one affine recogniser. + - **Multi-stream-var dependent** (``v: range(u() + w())`` referencing + more than one outer stream var): bbox is the affine combination over + both referents' ranges; guard threads through all dependencies. + - **Reverse-order dependent ranges**: e.g. ``v: range(u(), 0, -1)``; + needs to handle negative step and the corresponding reverse + enumeration. + """ + + @implements(Monoid.reduce) + def _(self, monoid: Monoid, body, streams: Streams): + stream_vars = set(streams.keys()) + + # streams of the form k: range(X) + simple_ranges = { + k: v + for (k, v) in streams.items() + if isinstance(v, Term) and _is_simple_range(v) + } + for u, u_stream in simple_ranges.items(): + if fvsof(u_stream) & stream_vars: + continue + + for v, v_stream in simple_ranges.items(): + if ( + isinstance(v_stream, Term) + and isinstance(_range_stop(v_stream), Term) + and _range_stop(v_stream).op == u + ): + fresh_streams = { + a: (u_stream if a == v else b) for (a, b) in streams.items() + } + + # there are other commuting rules for delta that we do not + # currently include + if isinstance(body, Term) and body.op == delta: + fresh_body = delta( + body.args[0], + jnp.where(v() < u(), body.args[1], monoid.identity), # type: ignore[arg-type] + ) + else: + fresh_body = jnp.where(v() < u(), body, monoid.identity) + + return monoid.reduce(fresh_body, fresh_streams) + + return fwd() + + +class ReduceRange(ObjectInterpretation): + """Replace concrete-range stream values with materialized ``jnp.arange``. + + reduce(M, streams ∪ {v: range(a, b, s)}, body) + ≡ reduce(M, streams ∪ {v: jnp.arange(a, b, s)}, body) + + when ``a``, ``b``, ``s`` are concrete and ``body`` is not a delta term. + Delegates the actual reduction to whichever handler picks up the + materialized ``jax.Array`` streams. + """ + + @implements(Monoid.reduce) + def _(self, monoid: Monoid, body, streams: Streams): + if isinstance(body, Term) and body.op == delta: + return fwd() + + new_streams: dict = {} + any_replaced = False + for k, v in streams.items(): + if isinstance(v, Term) and v.op == range: + new_streams[k] = jnp.arange( + _range_start(v), _range_stop(v), _range_step(v) + ) + any_replaced = True + else: + new_streams[k] = v + + if not any_replaced: + return fwd() + return monoid.reduce(body, new_streams) + + +# Cross-cutting delta rules not yet implemented: +# +# - **Delta-commuting** (DC-hoist): for any pure op ``f`` (no Scoped binders +# that intersect a delta's index ops), push delta outward: +# f(args..., delta(idx, body), args...) +# ≡ delta(idx, f(args..., body, args...)) +# This normalizes delta to the outermost position so the reduce rules can +# pattern-match ``isinstance(body, Term) and body.op == delta`` cleanly. +# The soundness condition is mechanical via ``op.__fvs_rule__``: refuse to +# commute when a non-delta arg's scope binds any op in the delta's idx. +# +# - **Delta-merging** (DC-merge): under a pure binary op ``f`` (or +# generalized n-ary), merge multiple deltas when their index tuples are +# subsequence-compatible: +# f(delta(idx_a, v), delta(idx_b, w)) ≡ delta(idx_max, f(v, w)) +# where ``idx_max`` is the longer of ``idx_a``, ``idx_b`` and ``idx_a`` is +# a subsequence of ``idx_b`` (or vice versa). Refuse to fire when neither +# is a subsequence of the other, since that would silently insert an +# outer-product broadcast. +# +# - **Empty-domain detection at the term level**: currently size-0 named +# dims must be resolved by leaf consumers (``bind_dims``, reductors with +# ``initial=monoid.identity``). The empty-domain check is intentionally +# NOT a rule on its own — rewrites stay size-polymorphic and leaf ops +# carry the burden. See the conversation in monoid.py's history for why. + + NormalizeIntp.extend( ArrayReduce(), + ReduceRange(), + ReduceDeltaIndependent(), + ReduceDependentRangeMask(), SumPlusJax(), ProductPlusJax(), MinPlusJax(), diff --git a/effectful/ops/monoid.py b/effectful/ops/monoid.py index 70bb50022..c9231510c 100644 --- a/effectful/ops/monoid.py +++ b/effectful/ops/monoid.py @@ -16,7 +16,6 @@ Scoped, deffn, implements, - iter_, syntactic_eq, syntactic_hash, ) @@ -96,8 +95,11 @@ def reduce[A, B, U: Body]( if isinstance(stream_body, Term): continue stream_values_iter = iter(stream_body) - if isinstance(stream_values_iter, Term) and stream_values_iter.op is iter_: + + # if we iterate and get a term instead of a real iterator, skip + if isinstance(stream_values_iter, Term): continue + new_reduces = [] for stream_val in stream_values_iter: with handler({stream_key: deffn(stream_val)}): diff --git a/tests/test_handlers_jax_monoid.py b/tests/test_handlers_jax_monoid.py index 35d041fe2..fe888ad43 100644 --- a/tests/test_handlers_jax_monoid.py +++ b/tests/test_handlers_jax_monoid.py @@ -1,11 +1,20 @@ import jax import pytest +from jax import random as random import effectful.handlers.jax.numpy as jnp from effectful.handlers.jax import bind_dims, unbind_dims -from effectful.handlers.jax.monoid import ArrayReduce, LogSumExp +from effectful.handlers.jax.monoid import ( + ArrayReduce, + LogSumExp, + ReduceDeltaIndependent, + ReduceDependentRangeMask, + delta, +) +from effectful.handlers.jax.monoid import range as Range from effectful.handlers.jax.scipy.special import logsumexp -from effectful.ops.monoid import Max, Min, Product, Sum +from effectful.ops.monoid import Max, Min, NormalizeIntp, Product, Sum +from effectful.ops.semantics import handler from tests._monoid_helpers import JAX_BACKEND, Backend, check_rewrite, define_vars MONOIDS = [ @@ -94,3 +103,167 @@ def test_reduce_array_3(monoid, reductor, backend: Backend): backend=backend, free_vars=[x, y, k1, k2, X, f, g], ) + + +# --------------------------------------------------------------------------- +# Delta rules. All tests use the operation form ``delta(idx, body)`` rather +# than the ``Delta`` dataclass; the delta op is the user-facing surface. +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("monoid,reductor", MONOIDS) +def test_reduce_delta_empty(monoid, reductor, backend: Backend): + """An empty-index delta unwraps to its body. + + reduce(M, streams, delta((), body)) ≡ reduce(M, streams, body) + """ + x = define_vars("x", typ=backend.scalar_typ) + X = define_vars("X", typ=backend.stream_typ) + + lhs = monoid.reduce(delta((), x()), {x: X()}) + rhs = monoid.reduce(x(), {x: X()}) + + check_rewrite( + lhs=lhs, + rhs=rhs, + rule=ReduceDeltaIndependent(), + backend=backend, + free_vars=[x, X], + ) + + +@pytest.mark.parametrize("monoid,reductor", MONOIDS) +def test_reduce_delta_independent_one(monoid, reductor, backend: Backend): + """One R1 step: peel the final preserved index off a delta. + + reduce(M, {y: Y()}, delta((y(),), f(y()))) + ≡ reduce(M, {}, delta((), bind_dims(f(unbind_dims(Y(), k)), k))) + """ + (y, k) = define_vars("y", "k", typ=backend.scalar_typ) + Y = define_vars("Y", typ=backend.stream_typ) + f = backend.fresh_op("f", n_args=1, ret="scalar") + + # We use a concrete range here instead of an abstract one, because + # unbind_dims is undefined on empty arrays (and the rewrite produces a + # different rhs in this case) + lhs = monoid.reduce(delta((y(),), f(y())), {y: Range(3)}) + rhs = monoid.reduce(bind_dims(f(unbind_dims(jnp.arange(3), k)), k), {}) + + check_rewrite( + lhs=lhs, + rhs=rhs, + rule=ReduceDeltaIndependent(), + backend=backend, + free_vars=[y, k, Y, f], + ) + + +@pytest.mark.parametrize("monoid,reductor", MONOIDS) +def test_reduce_delta_independent_preserves_others(monoid, reductor, backend: Backend): + """R1 peels only the final index. Streams not matching the peeled index op + stay untouched, as do earlier entries in the index tuple. + + reduce(M, {x: X(), y: Y()}, delta((x(), y()), f(x(), y()))) + ≡ reduce(M, {x: X()}, delta((x(),), bind_dims(f(x(), unbind_dims(Y(), k)), k))) + """ + (x, y, k) = define_vars("x", "y", "k", typ=backend.scalar_typ) + f = backend.fresh_op("f", n_args=2, ret="scalar") + + lhs = monoid.reduce(delta((x(), y()), f(x(), y())), {x: Range(2), y: Range(3)}) + rhs = monoid.reduce( + bind_dims( + bind_dims( + f(unbind_dims(jnp.arange(2), x), unbind_dims(jnp.arange(3), k)), k + ), + x, + ), + {}, + ) + + check_rewrite( + lhs=lhs, + rhs=rhs, + rule=ReduceDeltaIndependent(), + backend=backend, + free_vars=[f], + ) + + +@pytest.mark.parametrize("monoid,reductor", MONOIDS) +def test_reduce_dependent_range_mask(monoid, reductor, backend: Backend): + """A dependent range stream gets rewritten to the referent's bbox stream, + with the original constraint folded into the body as a where-guard. + + reduce(M, {u: range(0, N, 1), v: range(0, u(), 1)}, body) + ≡ reduce(M, {u: range(0, N, 1), v: range(0, N, 1)}, where(v() < u(), body, M.identity)) + """ + (u, v) = define_vars("u", "v", typ=backend.scalar_typ) + N = 5 + f = backend.fresh_op("f", n_args=2, ret="scalar") + + body = f(u(), v()) + + lhs = monoid.reduce(body, {u: Range(0, N, 1), v: Range(0, u(), 1)}) + rhs = monoid.reduce( + jnp.where(v() < u(), body, monoid.identity), + {u: Range(0, N, 1), v: Range(0, N, 1)}, + ) + + check_rewrite( + lhs=lhs, + rhs=rhs, + rule=ReduceDependentRangeMask(), + backend=backend, + free_vars=[u, v, f], + ) + + +@pytest.mark.parametrize("monoid,reductor", MONOIDS) +def test_reduce_dependent_range_mask_delta_body(monoid, reductor, backend: Backend): + """When the body is a delta term, R4 folds the constraint into the delta's + weight while leaving its index tuple untouched. + + reduce(M, {u: range(N), v: range(u())}, delta((u(), v()), w)) + ≡ reduce(M, {u: range(N), v: range(N)}, + delta((u(), v()), where(v() < u(), w, M.identity))) + """ + (u, v) = define_vars("u", "v", typ=backend.scalar_typ) + N = 5 + f = backend.fresh_op("f", n_args=2, ret="scalar") + + weight = f(u(), v()) + idx = (u(), v()) + + lhs = monoid.reduce(delta(idx, weight), {u: Range(0, N, 1), v: Range(0, u(), 1)}) + rhs = monoid.reduce( + delta(idx, jnp.where(v() < u(), weight, monoid.identity)), + {u: Range(0, N, 1), v: Range(0, N, 1)}, + ) + + check_rewrite( + lhs=lhs, + rhs=rhs, + rule=ReduceDependentRangeMask(), + backend=backend, + free_vars=[u, v, f], + ) + + +def test_reduce_matmul(): + key = jax.random.PRNGKey(0) + # Define dimensions + B, I, J, K = 2, 3, 4, 5 + + # Create sample matrices + X = random.normal(key, (B, I, J)) + Y = random.normal(key, (B, J, K)) + (b, i, j, k) = define_vars("b", "i", "j", "k", typ=jax.Array) + + with handler(NormalizeIntp): + actual = Sum.reduce( + delta((b(), i(), k()), unbind_dims(X, b, i, j) * unbind_dims(Y, b, j, k)), + {b: Range(B), i: Range(I), j: Range(J), k: Range(K)}, + ) + + expected = jnp.einsum("bij,bjk->bik", X, Y) + assert jnp.allclose(actual, expected) From 3a0a0aef7a91edd2835d271ec56b91d78aa2536c Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Wed, 17 Jun 2026 14:05:55 -0400 Subject: [PATCH 06/16] Add weighted streams (#665) * more precise stream type * add tests for weighted rules * add reduction rule for weighted streams and tests * add test to demo expectation * add numpyro monoid module * add quadrature * add tests * wip * refactor tests * wip * test composition of lifting and weighting * drop numpyro changes * drop unused ops * lint * make weighted a Monoid method * fix typing of jax arrays * change weighted typing to take callable * fix test * fix test * resolve type aliases before dispatching * wip * wip * remove typeof_full * wip * wip * wip * format * refactor test harness * drop unused test --- effectful/handlers/jax/_terms.py | 7 + effectful/handlers/jax/monoid.py | 15 +- effectful/ops/monoid.py | 125 ++++++- effectful/ops/semantics.py | 9 +- tests/_monoid_helpers.py | 466 +++++++++++++----------- tests/test_handlers_jax_monoid.py | 184 +++++----- tests/test_ops_monoid.py | 587 +++++++++++++++++------------- 7 files changed, 827 insertions(+), 566 deletions(-) diff --git a/effectful/handlers/jax/_terms.py b/effectful/handlers/jax/_terms.py index c88fe9341..05a5390e7 100644 --- a/effectful/handlers/jax/_terms.py +++ b/effectful/handlers/jax/_terms.py @@ -14,10 +14,17 @@ unbind_dims, ) from effectful.internals.tensor_utils import _desugar_tensor_index +from effectful.internals.unification import Box, nested_type from effectful.ops.syntax import defdata from effectful.ops.types import Expr, NotHandled, Operation, Term +@nested_type.register(jax.Array) +@nested_type.register(jax._src.core.Tracer) +def _(value): + return Box(jax.Array) + + class _IndexUpdateHelper: """Helper class to implement array-style .at[index].set() updates for effectful arrays.""" diff --git a/effectful/handlers/jax/monoid.py b/effectful/handlers/jax/monoid.py index 42d7866ec..3f6273be3 100644 --- a/effectful/handlers/jax/monoid.py +++ b/effectful/handlers/jax/monoid.py @@ -16,6 +16,7 @@ Product, Streams, Sum, + distributes_over, outer_stream, ) from effectful.ops.semantics import evaluate, fvsof, fwd, handler, typeof @@ -38,6 +39,10 @@ def cartesian_prod(x, y): LogSumExp = Monoid(name="LogSumExp", identity=jnp.asarray(float("-inf"))) +# ``Sum`` in log space is multiplication, which distributes over ``LogSumExp``: +# a + logsumexp(b, c) = logsumexp(a + b, a + c) +distributes_over.register(Sum, LogSumExp) + def _jax_args(args): """True iff ``args`` is non-empty and every arg is a concrete @@ -108,7 +113,15 @@ def plus(self, *args): if not isinstance(a, jax.Array): return fwd() result = a if result is None else cartesian_prod(result, a) - return result if result is not None else CartesianProduct.identity + if result is None: + return CartesianProduct.identity + # CartesianProduct values are streams of rows. ``cartesian_prod`` + # already lifts 1D inputs to 2D, but a single-array call seeds + # ``result = a`` unchanged — promote so the rank invariant holds for + # every array-path return. + if result.ndim == 1: + result = result[:, None] + return result ARRAY_REDUCTORS = { diff --git a/effectful/ops/monoid.py b/effectful/ops/monoid.py index c9231510c..76351fa62 100644 --- a/effectful/ops/monoid.py +++ b/effectful/ops/monoid.py @@ -10,7 +10,14 @@ from typing import Annotated, Any from effectful.internals.disjoint_set import DisjointSet -from effectful.ops.semantics import coproduct, evaluate, fvsof, fwd, handler +from effectful.ops.semantics import ( + coproduct, + evaluate, + fvsof, + fwd, + handler, + typeof, +) from effectful.ops.syntax import ( ObjectInterpretation, Scoped, @@ -19,11 +26,17 @@ syntactic_eq, syntactic_hash, ) -from effectful.ops.types import Expr, Interpretation, NotHandled, Operation, Term +from effectful.ops.types import ( + Expr, + Interpretation, + NotHandled, + Operation, + Term, +) + +type Stream[T] = Iterable[T] -# Note: The streams value type should be something like Iterable[T], but some of -# our target stream types (e.g. jax.Array) are not subtypes of Iterable -type Streams[T] = Mapping[Operation[[], T], Any] +type Streams = Mapping[Operation[[], Any], Stream[Any]] type Body[T] = ( Iterable[T] @@ -34,9 +47,7 @@ ) -def outer_stream( - streams: Streams, -) -> Iterable[tuple[Operation, Expr, dict[Operation, Expr]]]: +def outer_stream(streams: Streams) -> Iterable[tuple[Operation, Stream, Streams]]: """Returns the streams that can be ordered outermost in the loop nest as well as the remaining streams in the nest. @@ -51,13 +62,13 @@ def outer_stream( ) -class Monoid[T]: +class Monoid[W]: """A monoid with ``plus`` and ``reduce`` :class:`Operation` s.""" _name: str - identity: T + identity: W - def __init__(self, identity: T, name: str): + def __init__(self, identity: W, name: str): self._name = name self.identity = identity @@ -111,6 +122,18 @@ def reduce[A, B, U: Body]( return self.plus(*new_reduces) raise NotHandled + @Operation.define + def weighted[T]( + self, stream: Stream[T], weight: Callable[[T], W] | Operation[[T], W] + ) -> Stream[T]: + """A stream paired with a per-element weight. ``var`` is an + :class:`Operation` standing for "an element of ``stream``"; ``weight`` + is an expression that uses ``var`` and evaluates to the weight of that + element. + + """ + raise NotHandled + class MonoidWithZero[T](Monoid[T]): zero: T @@ -175,6 +198,12 @@ def _is_monoid_reduce(op: Operation) -> bool: return isinstance(owner, Monoid) and op is owner.reduce +def _is_monoid_weighted(op: Operation) -> bool: + """True if ``op`` is the ``weighted`` operation of some :class:`Monoid`.""" + owner = getattr(op, "__self__", None) + return isinstance(owner, Monoid) and op is owner.weighted + + class PlusEmpty(ObjectInterpretation): """plus() = 0""" @@ -557,6 +586,78 @@ def reduce(self, sum_monoid: Monoid, sum_body, sum_streams): return fwd() +class ReduceWeightedStream(ObjectInterpretation): + """reduce(M, body, {x: WM.weighted(s, v, w), ...}) = reduce(M, WM.plus(w[v:=x()], body), {x: s, ...}) + + requires distributes_over(WM, M). + + The substitution ``v -> x`` is done by beta-reducing ``deffn(w, v)`` on + ``x()`` — symbolic, no Python dispatch on the weight expression. + """ + + @implements(Monoid.reduce) + def reduce(self, monoid, body, streams): + for k, v in streams.items(): + if isinstance(v, Term) and _is_monoid_weighted(v.op): + v_stream, v_weight = v.args + v_monoid = v.op.__self__ + if not distributes_over(v_monoid, monoid): + continue + w_at_k = v_weight(k()) + weighted_body = v_monoid.plus(w_at_k, body) + new_streams = {**streams, k: v_stream} + return monoid.reduce(weighted_body, new_streams) + return fwd() + + +class ReduceCartesianWeightedStream(ObjectInterpretation): + """``CartesianProduct.reduce`` over a :func:`weighted` body whose + ``weight`` is independent of the plate (product-index) streams:: + + CartesianProduct.reduce(M.weighted(s, w), plates) + = M.weighted( + CartesianProduct.reduce(s, plates), + deffn(M.reduce(w, {e: row()}), row), + ) + + Reuses ``body``'s element binder ``e`` (already typed by construction); + introduces a fresh ``row`` binder typed as ``Iterable[elem_type]``. + + Only fires when ``w`` is independent of the plate vars. + """ + + @Operation.define + @staticmethod + def _iterable_elem[T](iter: Iterable[T]) -> T: + raise NotHandled + + @implements(Monoid.reduce) + def reduce(self, monoid, body, streams): + if monoid is not CartesianProduct: + return fwd() + if not (isinstance(body, Term) and _is_monoid_weighted(body.op)): + return fwd() + + s, w = body.args + if not isinstance(s, Term) and len(s) == 0: + return CartesianProduct.reduce([], streams) + + if set(streams.keys()) & fvsof(w): + return fwd() + + elem_typ = typeof(self._iterable_elem(s)) + elem_op = Operation.define(elem_typ, name="elem") + row_op = Operation.define(Iterable[elem_typ], name="row") + + weight_monoid = body.op.__self__ + joint_weight = deffn( + weight_monoid.reduce(w(elem_op()), {elem_op: row_op()}), row_op + ) + joint_stream = CartesianProduct.reduce(s, streams) + + return weight_monoid.weighted(joint_stream, joint_weight) + + class MonoidOverCallable(ObjectInterpretation): """``monoid.reduce(f, streams) = lambda *a: monoid.reduce(f(*a), streams)``.""" @@ -751,6 +852,8 @@ def extend(self, *intps: Interpretation) -> typing.Self: ReduceSplit(), ReduceFactorization(), ReduceDistributeCartesianProduct(), + ReduceWeightedStream(), + ReduceCartesianWeightedStream(), PlusEmpty(), PlusSingle(), PlusIdentity(), diff --git a/effectful/ops/semantics.py b/effectful/ops/semantics.py index 16ace787a..8937a3ae5 100644 --- a/effectful/ops/semantics.py +++ b/effectful/ops/semantics.py @@ -301,6 +301,13 @@ def _evaluate_list_view(expr, **kwargs): def _simple_type(tp: type) -> type: """Convert a type object into a type that can be dispatched on.""" + + def _resolve_aliases(tp: type) -> type: + tp = typing.get_origin(tp) or tp + if isinstance(tp, typing.TypeAliasType): + return _resolve_aliases(tp.__value__) + return tp + if isinstance(tp, typing.TypeVar): tp = ( tp.__bound__ @@ -318,7 +325,7 @@ def _simple_type(tp: type) -> type: tp = functools.reduce(operator.or_, (type(arg) for arg in args)) if isinstance(tp, types.UnionType): raise TypeError(f"Union types are not supported: {tp}") - return typing.get_origin(tp) or tp + return _resolve_aliases(tp) class _TypeofIntp(ObjectInterpretation): diff --git a/tests/_monoid_helpers.py b/tests/_monoid_helpers.py index f15103e30..f8089bec7 100644 --- a/tests/_monoid_helpers.py +++ b/tests/_monoid_helpers.py @@ -1,146 +1,22 @@ +import builtins import itertools -from collections.abc import Callable, Mapping, Sequence -from dataclasses import dataclass -from typing import Any, get_args, get_origin +import typing +from abc import ABC, abstractmethod +from collections.abc import Callable, Mapping +from typing import Any, Literal, overload import jax from hypothesis import given, settings from hypothesis import strategies as st +from hypothesis.strategies import SearchStrategy import effectful.handlers.jax.numpy as _jnp from effectful.internals.runtime import interpreter -from effectful.ops.monoid import NormalizeIntp -from effectful.ops.semantics import apply, evaluate, handler +from effectful.ops.monoid import NormalizeIntp, Stream, _is_monoid_weighted +from effectful.ops.semantics import apply, evaluate, fvsof, handler from effectful.ops.syntax import _BaseTerm, defdata, deffn, syntactic_eq from effectful.ops.types import NotHandled, Operation, Term -_JAX_ARRAY_SHAPE = (2,) - - -def _jax_array_value_strategy() -> st.SearchStrategy[jax.Array]: - return st.lists( - st.integers(min_value=-5, max_value=5), - min_size=_JAX_ARRAY_SHAPE[0], - max_size=_JAX_ARRAY_SHAPE[0], - ).map(lambda xs: jax.numpy.asarray(xs, dtype=jax.numpy.float32)) - - -# Unary jax fns map a scalar to a 1-D array (analogous to ``_UNARY_LIST_FNS`` -# for ints). Uses the effectful-wrapped jnp so named-dim broadcasting works. -_UNARY_JAX_FNS: list[Callable[[jax.Array], jax.Array]] = [ - lambda a: _jnp.stack([a, a + 1]), - lambda a: _jnp.stack([a, -a]), - lambda a: _jnp.stack([a, a + 1, 2 * a]), -] - -_BINARY_JAX_FNS: list[Callable[[jax.Array, jax.Array], jax.Array]] = [ - lambda a, b: a + b, - lambda a, b: a - b, - lambda a, b: a * b, -] - - -def _value_strategy_for(annotation: Any) -> st.SearchStrategy[Any]: - """Strategy for the value an *0-arg* Operation should return.""" - if annotation is int: - return st.integers(min_value=-100, max_value=100) - if annotation is float: - return st.floats(allow_nan=False) - if get_origin(annotation) is list and get_args(annotation) == (int,): - return st.lists(st.integers(min_value=-100, max_value=100), max_size=2) - if annotation is jax.Array: - return _jax_array_value_strategy() - if get_origin(annotation) is list and get_args(annotation) == (jax.Array,): - return st.lists(_jax_array_value_strategy(), max_size=2) - raise NotImplementedError( - f"No value strategy for return annotation {annotation!r}; " - "supported: int, list[int], jax.Array, list[jax.Array]" - ) - - -_UNARY_NUM_FNS: list[Callable[[int], int]] = [ - lambda x: x, - lambda x: x + 1, - lambda x: x - 1, - lambda x: -x, - lambda x: 2 * x, - lambda x: 3 * x + 1, -] - -_BINARY_NUM_FNS: list[Callable[[int, int], int]] = [ - lambda x, y: x + y, - lambda x, y: x - y, - lambda x, y: x * y, - lambda x, y: x + 2 * y, - lambda x, y: 2 * x - y, -] - -_UNARY_LIST_FNS: list[Callable[[int], list[int]]] = [ - lambda _x: [], - lambda x: [x], - lambda x: [x, x + 1], - lambda x: [x, -x], - lambda x: [0, x, x + 1], -] - -_UNARY_JAX_LIST_FNS: list[Callable[[jax.Array], list[jax.Array]]] = [ - lambda _x: [], - lambda x: [x], - lambda x: [x, x + 1], - lambda x: [x, -x], -] - - -def _strategy_for_op(op: Operation) -> st.SearchStrategy[Callable[..., Any]]: - """Pick a strategy producing a callable suitable for binding `op` in an - interpretation. Inspects the operation's signature. - """ - sig = op.__signature__ - params = list(sig.parameters.values()) - ret = sig.return_annotation - param_types = tuple(p.annotation for p in params) - - if not params: - return _value_strategy_for(ret).map(deffn) - if ret in (int, float) and param_types == (int,): - return st.sampled_from(_UNARY_NUM_FNS) - if ret in (int, float) and param_types == (int, int): - return st.sampled_from(_BINARY_NUM_FNS) - if get_origin(ret) is list and get_args(ret) == (int,) and param_types == (int,): - return st.sampled_from(_UNARY_LIST_FNS) - if ret is jax.Array and param_types == (jax.Array,): - return st.sampled_from(_UNARY_JAX_FNS) - if ret is jax.Array and param_types == (jax.Array, jax.Array): - return st.sampled_from(_BINARY_JAX_FNS) - if ( - get_origin(ret) is list - and get_args(ret) == (jax.Array,) - and param_types == (jax.Array,) - ): - return st.sampled_from(_UNARY_JAX_LIST_FNS) - raise NotImplementedError( - f"No callable strategy for free var with return {ret!r}, params {param_types!r}" - ) - - -@st.composite -def random_interpretation( - draw: st.DrawFn, free_vars: Sequence[Operation] -) -> Mapping[Operation, Callable[..., Any]]: - """Draw an Interpretation binding every Operation in `case.free_vars` to - a randomly chosen value/callable. Keys are Operation identities. - """ - intp: dict[Operation, Callable[..., Any]] = {} - for op in free_vars: - intp[op] = draw(_strategy_for_op(op)) - return intp - - -def define_vars(*names, typ=int): - if len(names) == 1: - return Operation.define(typ, name=names[0]) - return tuple(Operation.define(typ, name=n) for n in names) - def syntactic_eq_alpha(x, y) -> bool: """Alpha-equivalence-respecting variant of ``syntactic_eq``. @@ -251,8 +127,7 @@ def _apply_canonical(op, *args, **kwargs) -> Term: return evaluate(expr) -@dataclass(frozen=True) -class Backend: +class Backend(ABC): """A value-domain spec used to share monoid tests across int and jax.Array backends. Provides the concrete value type, the hypothesis strategy for drawing scalars in property tests, and an equality predicate that works @@ -262,10 +137,29 @@ class Backend: name: str scalar_typ: Any stream_typ: Any - scalar_strategy: st.SearchStrategy[Any] - eq: Callable[[Any, Any], bool] - - def fresh_op(self, name: str, n_args: int = 1, ret: str = "scalar") -> Operation: + strategy_for_op: dict[Operation, st.SearchStrategy[Callable[..., Any]]] + + def __init__(self): + self.strategy_for_op = {} + + @abstractmethod + 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", + ) -> Operation: """Build a fresh, unhandled Operation whose parameter and return annotations are derived from this backend. @@ -275,81 +169,245 @@ def fresh_op(self, name: str, n_args: int = 1, ret: str = "scalar") -> Operation """ scalar = self.scalar_typ out = self.stream_typ if ret == "stream" else scalar - params = ", ".join(f"_a{i}" for i in range(n_args)) + 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) fn = ns["_fn"] fn.__annotations__ = { - **{f"_a{i}": scalar for i in range(n_args)}, + **{f"_a{i}": t for i, t in enumerate(arg_types)}, "return": out, } - return Operation.define(fn, name=name) + op = Operation.define(fn, name=name) + self.strategy_for_op[op] = self.strategy(arg_types, ret) + return op + + @overload + def define_vars(self, name: str, /, **kwargs) -> Operation: ... + + @overload + def define_vars( + self, n1: str, n2: str, /, *names: str, **kwargs + ) -> tuple[Operation, ...]: ... + + def define_vars(self, *names: str, **kwargs) -> Operation | tuple[Operation, ...]: # type: ignore[misc] + if len(names) == 1: + return self._fresh_op(names[0], **kwargs) + return tuple(self._fresh_op(n, **kwargs) for n in names) + + def check_rewrite( + self, + lhs, + rhs, + rule, + *, + max_examples: int = 25, + deadline=None, + normalize=NormalizeIntp, + ) -> None: + with handler(rule): + norm = evaluate(lhs) + assert syntactic_eq_alpha(norm, rhs) + + fvs = fvsof(lhs) | fvsof(rhs) + + @st.composite + def random_interpretation( + draw: st.DrawFn, + ) -> Mapping[Operation, Callable[..., Any]]: + """Draw an Interpretation binding every Operation in `free_vars` to + a randomly chosen value/callable. Keys are Operation identities. + """ + intp: dict[Operation, Callable[..., Any]] = {} + for op, strategy in self.strategy_for_op.items(): + if op in fvs: + intp[op] = draw(strategy) + return intp + + @given(intp=random_interpretation()) + @settings( + max_examples=max_examples, deadline=deadline, report_multiple_bugs=False + ) + def _check_semantics(intp): + with handler(normalize), handler(intp): + lhs_val = evaluate(lhs) + rhs_val = evaluate(rhs) + assert self.eq(lhs_val, rhs_val) + _check_semantics() -def _int_eq(a: Any, b: Any) -> bool: - return not isinstance(a, Term) and not isinstance(b, Term) and a == b +def _is_weighted(x: Any) -> bool: + return isinstance(x, Term) and _is_monoid_weighted(x.op) -def _jax_eq(a: Any, b: Any) -> bool: - def _leaf_eq(x: Any, y: Any) -> bool: - return bool(jax.numpy.all(jax.numpy.isclose(x, y, equal_nan=True))) - try: - leaves = jax.tree.leaves(jax.tree.map(_leaf_eq, a, b)) - except (ValueError, TypeError): +def _weight_pairs(x: Any, monoid: Any) -> list[tuple[Any, Any]] | None: + """Return ``(element, weight)`` pairs for a stream. + + A weighted-monoid Term yields each element paired with its weight. A plain + (unweighted) stream yields each element paired with ``monoid.identity`` -- + the no-op weight -- so an unweighted stream compares equal to a weighted one + exactly when every weight reduces to the identity (e.g. ``[()]`` vs a + weighted ``[()]`` whose single empty row reduces to the identity, and, more + generally, whenever both streams are empty). Returns ``None`` for a + non-stream Term, which never compares equal to a weighted stream. + """ + if isinstance(x, Term): + if not _is_monoid_weighted(x.op): + return None + stream, weight = x.args + assert not isinstance(stream, Term) + return [(e, typing.cast(Callable, weight)(e)) for e in stream] + return [(e, monoid.identity) for e in x] + + +def _weighted_stream_eq(a, b, leaf_eq: Callable[[Any, Any], bool]) -> bool: + monoids = {x.op.__self__ for x in (a, b) if _is_weighted(x)} + # distinct weight monoids can never be equal + if len(monoids) != 1: return False - return all(leaves) - - -def check_rewrite( - lhs, - rhs, - rule, - *, - backend: Backend, - free_vars=[], - max_examples: int = 25, - deadline=None, -) -> None: - with handler(rule): - norm = evaluate(lhs) - assert syntactic_eq_alpha(norm, rhs) - - @given(intp=random_interpretation(free_vars)) - @settings(max_examples=max_examples, deadline=deadline) - def _check_semantics(intp): - with handler(NormalizeIntp), handler(intp): - lhs_val = evaluate(lhs) - rhs_val = evaluate(rhs) - assert backend.eq(lhs_val, rhs_val) - - _check_semantics() - - -INT_BACKEND = Backend( - name="int", - scalar_typ=int, - stream_typ=list[int], - scalar_strategy=st.integers(min_value=-100, max_value=100), - eq=_int_eq, -) - - -JAX_BACKEND = Backend( - name="jax", - scalar_typ=jax.Array, - stream_typ=jax.Array, - scalar_strategy=_jax_array_value_strategy(), - eq=_jax_eq, -) - - -__all__ = [ - "Backend", - "INT_BACKEND", - "JAX_BACKEND", - "random_interpretation", - "define_vars", - "syntactic_eq_alpha", - "check_rewrite", -] + monoid = next(iter(monoids)) + + a_pairs = _weight_pairs(a, monoid) + b_pairs = _weight_pairs(b, monoid) + if a_pairs is None or b_pairs is None or len(a_pairs) != len(b_pairs): + return False + for (ea, wa), (eb, wb) in zip(a_pairs, b_pairs): + if not leaf_eq(ea, eb) or not leaf_eq(wa, wb): + return False + return True + + +class IntBackend(Backend): + name = "int" + scalar_typ = int + stream_typ = Stream[int] + + _unary_num_fns: list[Callable[[int], int]] = [ + lambda x: x, + lambda x: x + 1, + lambda x: x - 1, + lambda x: -x, + lambda x: 2 * x, + lambda x: 3 * x + 1, + ] + + _binary_num_fns: list[Callable[[int, int], int]] = [ + lambda x, y: x + y, + lambda x, y: x - y, + lambda x, y: x * y, + lambda x, y: x + 2 * y, + lambda x, y: 2 * x - y, + ] + + _unary_list_fns: list[Callable[[int], list[int]]] = [ + lambda _x: [], + lambda x: [x], + lambda x: [x, x + 1], + lambda x: [x, -x], + lambda x: [0, x, x + 1], + ] + + def strategy( + self, + arg_types: tuple[type, ...] = (), + ret: Literal["scalar", "stream"] = "scalar", + ) -> SearchStrategy: + match arg_types, ret: + case (), "scalar": + 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) + case (builtins.int, builtins.int), "scalar": + return st.sampled_from(self._binary_num_fns) + case (builtins.int,), "stream": + return st.sampled_from(self._unary_list_fns) + raise NotImplementedError( + f"No int strategy for op with return {ret!r} and {arg_types} args" + ) + + def eq(self, a: Any, b: Any) -> bool: + if _is_weighted(a) or _is_weighted(b): + return _weighted_stream_eq(a, b, self.eq) + return not isinstance(a, Term) and not isinstance(b, Term) and a == b + + +class JaxBackend(Backend): + name = "jax" + scalar_typ = jax.Array + stream_typ = jax.Array + + _unary_jax_scalar_fns: list[Callable[[jax.Array], jax.Array]] = [ + lambda a: a, + lambda a: a + 1, + lambda a: a - 1, + lambda a: -a, + lambda a: 2 * a, + ] + + _unary_jax_stream_fns: list[Callable[[jax.Array], Stream[jax.Array]]] = [ + lambda a: _jnp.stack([a, a + 1]), + lambda a: _jnp.stack([a, -a]), + lambda a: _jnp.stack([a, a + 1, 2 * a]), + ] + + _binary_jax_scalar_fns: list[Callable[[jax.Array, jax.Array], jax.Array]] = [ + lambda a, b: a + b, + lambda a, b: a - b, + lambda a, b: a * b, + ] + + def strategy( + self, + arg_types: tuple[type, ...] = (), + ret: Literal["scalar", "stream"] = "scalar", + ) -> st.SearchStrategy[Callable]: + match arg_types, ret: + case (), "scalar": + return ( + st.lists( + st.integers(min_value=-5, max_value=5), + min_size=2, + max_size=2, + ) + .map(lambda xs: jax.numpy.asarray(xs, dtype=jax.numpy.float32)) + .map(deffn) + ) + case (), "stream": + return ( + st.lists( + st.integers(min_value=-5, max_value=5), + min_size=1, + max_size=2, + ) + .map(lambda xs: jax.numpy.asarray(xs, dtype=jax.numpy.float32)) + .map(deffn) + ) + case (jax.Array,), "scalar": + return st.sampled_from(self._unary_jax_scalar_fns) + case (jax.Array, jax.Array), "scalar": + return st.sampled_from(self._binary_jax_scalar_fns) + case (jax.Array,), "stream": + return st.sampled_from(self._unary_jax_stream_fns) + + raise NotImplementedError( + f"No jax strategy for op with return {ret!r} and {arg_types} args" + ) + + def eq(self, a: Any, b: Any) -> bool: + if _is_weighted(a) or _is_weighted(b): + return _weighted_stream_eq(a, b, self.eq) + + def _leaf_eq(x: Any, y: Any) -> bool: + return bool(jax.numpy.all(jax.numpy.isclose(x, y, equal_nan=True))) + + try: + leaves = jax.tree.leaves(jax.tree.map(_leaf_eq, a, b)) + except (ValueError, TypeError): + return False + return all(leaves) + + +__all__ = ["Backend", "IntBackend", "JaxBackend", "syntactic_eq_alpha"] diff --git a/tests/test_handlers_jax_monoid.py b/tests/test_handlers_jax_monoid.py index fe888ad43..18df84018 100644 --- a/tests/test_handlers_jax_monoid.py +++ b/tests/test_handlers_jax_monoid.py @@ -1,3 +1,6 @@ +import functools +import typing + import jax import pytest from jax import random as random @@ -7,15 +10,24 @@ from effectful.handlers.jax.monoid import ( ArrayReduce, LogSumExp, + ProductPlusJax, ReduceDeltaIndependent, ReduceDependentRangeMask, delta, ) from effectful.handlers.jax.monoid import range as Range from effectful.handlers.jax.scipy.special import logsumexp -from effectful.ops.monoid import Max, Min, NormalizeIntp, Product, Sum -from effectful.ops.semantics import handler -from tests._monoid_helpers import JAX_BACKEND, Backend, check_rewrite, define_vars +from effectful.ops.monoid import ( + Max, + Min, + NormalizeIntp, + Product, + ReduceWeightedStream, + Sum, +) +from effectful.ops.semantics import coproduct, handler +from effectful.ops.types import Interpretation +from tests._monoid_helpers import JaxBackend MONOIDS = [ pytest.param(Sum, jnp.sum, id="Sum"), @@ -27,28 +39,27 @@ @pytest.fixture -def backend() -> Backend: - return JAX_BACKEND +def backend() -> JaxBackend: + return JaxBackend() @pytest.mark.parametrize("monoid,reductor", MONOIDS) -def test_reduce_array_1(monoid, reductor, backend: Backend): - (x, k) = define_vars("x", "k", typ=jax.Array) - X = define_vars("X", typ=backend.stream_typ) +def test_reduce_array_1(monoid, reductor, backend: JaxBackend): + (x, k) = backend.define_vars("x", "k", ret="scalar") + X = backend.define_vars("X", ret="stream") lhs = monoid.reduce(x(), {x: X()}) rhs = reductor(bind_dims(unbind_dims(X(), k), k), axis=0) - - check_rewrite( - lhs=lhs, rhs=rhs, rule=ArrayReduce(), backend=backend, free_vars=[x, X, k] - ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ArrayReduce()) @pytest.mark.parametrize("monoid,reductor", MONOIDS) -def test_reduce_array_2(monoid, reductor, backend: Backend): - (x, y, k1, k2) = define_vars("x", "y", "k1", "k2", typ=backend.scalar_typ) - (X, Y) = define_vars("X", "Y", typ=backend.stream_typ) - f = backend.fresh_op("f", n_args=2, ret="scalar") +def test_reduce_array_2(monoid, reductor, backend: JaxBackend): + (x, y, k1, k2) = backend.define_vars("x", "y", "k1", "k2", ret="scalar") + (X, Y) = backend.define_vars("X", "Y", ret="stream") + f = backend.define_vars( + "f", arg_types=(backend.scalar_typ, backend.scalar_typ), ret="scalar" + ) lhs = monoid.reduce(f(x(), y()), {x: X(), y: Y()}) rhs = reductor( @@ -61,25 +72,20 @@ def test_reduce_array_2(monoid, reductor, backend: Backend): ), axis=0, ) - - check_rewrite( - lhs=lhs, - rhs=rhs, - rule=ArrayReduce(), - backend=backend, - free_vars=[x, y, k1, k2, X, Y, f], - ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ArrayReduce()) @pytest.mark.parametrize("monoid,reductor", MONOIDS) -def test_reduce_array_3(monoid, reductor, backend: Backend): +def test_reduce_array_3(monoid, reductor, backend: JaxBackend): """Stream `y` is `g(x())` — depends on the bound element of X. The reducer must inline ``g`` along the same named dim used to unbind `x`.""" - (x, y, k1, k2) = define_vars("x", "y", "k1", "k2", typ=backend.scalar_typ) - X = define_vars("X", typ=backend.stream_typ) + (x, y, k1, k2) = backend.define_vars("x", "y", "k1", "k2", ret="scalar") + X = backend.define_vars("X", ret="stream") - f = backend.fresh_op("f", n_args=2, ret="scalar") - g = backend.fresh_op("g", n_args=1, ret="stream") + f = backend.define_vars( + "f", arg_types=[backend.scalar_typ, backend.scalar_typ], ret="scalar" + ) + g = backend.define_vars("g", arg_types=[backend.scalar_typ], ret="stream") lhs = monoid.reduce(f(x(), y()), {x: X(), y: g(x())}) rhs = reductor( @@ -95,13 +101,37 @@ def test_reduce_array_3(monoid, reductor, backend: Backend): ), axis=0, ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ArrayReduce()) - check_rewrite( + +def test_jax_weighted_reduce(backend: JaxBackend): + """Sum over a single stream with ``Product`` weights lowers to + ``jnp.sum(w(X) * body(X))`` under ``NormalizeIntp`` ∘ ``ArrayReduce``. + + Verifies that the desugaring rule composes cleanly with the JAX lowering + so existing handlers need no changes to support weighted streams. + + """ + (x, k) = backend.define_vars("x", "k", ret="scalar") + X = backend.define_vars("X", ret="stream") + body = backend.define_vars("body", arg_types=[backend.scalar_typ], ret="scalar") + w = backend.define_vars("w", arg_types=[backend.scalar_typ], ret="scalar") + + ws = Product.weighted(X(), w) + lhs = Sum.reduce(body(x()), {x: ws}) + rhs = jnp.sum( + bind_dims(w(unbind_dims(X(), k)) * body(unbind_dims(X(), k)), k), axis=0 + ) + backend.check_rewrite( lhs=lhs, rhs=rhs, - rule=ArrayReduce(), - backend=backend, - free_vars=[x, y, k1, k2, X, f, g], + rule=functools.reduce( + coproduct, + typing.cast( + list[Interpretation], + [ReduceWeightedStream(), ArrayReduce(), ProductPlusJax()], + ), + ), ) @@ -112,62 +142,51 @@ def test_reduce_array_3(monoid, reductor, backend: Backend): @pytest.mark.parametrize("monoid,reductor", MONOIDS) -def test_reduce_delta_empty(monoid, reductor, backend: Backend): +def test_reduce_delta_empty(monoid, reductor, backend: JaxBackend): """An empty-index delta unwraps to its body. reduce(M, streams, delta((), body)) ≡ reduce(M, streams, body) """ - x = define_vars("x", typ=backend.scalar_typ) - X = define_vars("X", typ=backend.stream_typ) + x = backend.define_vars("x", ret="scalar") + X = backend.define_vars("X", ret="stream") lhs = monoid.reduce(delta((), x()), {x: X()}) rhs = monoid.reduce(x(), {x: X()}) - - check_rewrite( - lhs=lhs, - rhs=rhs, - rule=ReduceDeltaIndependent(), - backend=backend, - free_vars=[x, X], - ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceDeltaIndependent()) @pytest.mark.parametrize("monoid,reductor", MONOIDS) -def test_reduce_delta_independent_one(monoid, reductor, backend: Backend): +def test_reduce_delta_independent_one(monoid, reductor, backend: JaxBackend): """One R1 step: peel the final preserved index off a delta. reduce(M, {y: Y()}, delta((y(),), f(y()))) ≡ reduce(M, {}, delta((), bind_dims(f(unbind_dims(Y(), k)), k))) """ - (y, k) = define_vars("y", "k", typ=backend.scalar_typ) - Y = define_vars("Y", typ=backend.stream_typ) - f = backend.fresh_op("f", n_args=1, ret="scalar") + (y, k) = backend.define_vars("y", "k", ret="scalar") + f = backend.define_vars("f", arg_types=[backend.scalar_typ], ret="scalar") # We use a concrete range here instead of an abstract one, because # unbind_dims is undefined on empty arrays (and the rewrite produces a # different rhs in this case) lhs = monoid.reduce(delta((y(),), f(y())), {y: Range(3)}) rhs = monoid.reduce(bind_dims(f(unbind_dims(jnp.arange(3), k)), k), {}) - - check_rewrite( - lhs=lhs, - rhs=rhs, - rule=ReduceDeltaIndependent(), - backend=backend, - free_vars=[y, k, Y, f], - ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceDeltaIndependent()) @pytest.mark.parametrize("monoid,reductor", MONOIDS) -def test_reduce_delta_independent_preserves_others(monoid, reductor, backend: Backend): +def test_reduce_delta_independent_preserves_others( + monoid, reductor, backend: JaxBackend +): """R1 peels only the final index. Streams not matching the peeled index op stay untouched, as do earlier entries in the index tuple. reduce(M, {x: X(), y: Y()}, delta((x(), y()), f(x(), y()))) ≡ reduce(M, {x: X()}, delta((x(),), bind_dims(f(x(), unbind_dims(Y(), k)), k))) """ - (x, y, k) = define_vars("x", "y", "k", typ=backend.scalar_typ) - f = backend.fresh_op("f", n_args=2, ret="scalar") + (x, y, k) = backend.define_vars("x", "y", "k", ret="scalar") + f = backend.define_vars( + "f", arg_types=[backend.scalar_typ, backend.scalar_typ], ret="scalar" + ) lhs = monoid.reduce(delta((x(), y()), f(x(), y())), {x: Range(2), y: Range(3)}) rhs = monoid.reduce( @@ -179,27 +198,22 @@ def test_reduce_delta_independent_preserves_others(monoid, reductor, backend: Ba ), {}, ) - - check_rewrite( - lhs=lhs, - rhs=rhs, - rule=ReduceDeltaIndependent(), - backend=backend, - free_vars=[f], - ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceDeltaIndependent()) @pytest.mark.parametrize("monoid,reductor", MONOIDS) -def test_reduce_dependent_range_mask(monoid, reductor, backend: Backend): +def test_reduce_dependent_range_mask(monoid, reductor, backend: JaxBackend): """A dependent range stream gets rewritten to the referent's bbox stream, with the original constraint folded into the body as a where-guard. reduce(M, {u: range(0, N, 1), v: range(0, u(), 1)}, body) ≡ reduce(M, {u: range(0, N, 1), v: range(0, N, 1)}, where(v() < u(), body, M.identity)) """ - (u, v) = define_vars("u", "v", typ=backend.scalar_typ) + (u, v) = backend.define_vars("u", "v", ret="scalar") N = 5 - f = backend.fresh_op("f", n_args=2, ret="scalar") + f = backend.define_vars( + "f", arg_types=[backend.scalar_typ, backend.scalar_typ], ret="scalar" + ) body = f(u(), v()) @@ -208,18 +222,11 @@ def test_reduce_dependent_range_mask(monoid, reductor, backend: Backend): jnp.where(v() < u(), body, monoid.identity), {u: Range(0, N, 1), v: Range(0, N, 1)}, ) - - check_rewrite( - lhs=lhs, - rhs=rhs, - rule=ReduceDependentRangeMask(), - backend=backend, - free_vars=[u, v, f], - ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceDependentRangeMask()) @pytest.mark.parametrize("monoid,reductor", MONOIDS) -def test_reduce_dependent_range_mask_delta_body(monoid, reductor, backend: Backend): +def test_reduce_dependent_range_mask_delta_body(monoid, reductor, backend: JaxBackend): """When the body is a delta term, R4 folds the constraint into the delta's weight while leaving its index tuple untouched. @@ -227,9 +234,11 @@ def test_reduce_dependent_range_mask_delta_body(monoid, reductor, backend: Backe ≡ reduce(M, {u: range(N), v: range(N)}, delta((u(), v()), where(v() < u(), w, M.identity))) """ - (u, v) = define_vars("u", "v", typ=backend.scalar_typ) + (u, v) = backend.define_vars("u", "v", ret="scalar") N = 5 - f = backend.fresh_op("f", n_args=2, ret="scalar") + f = backend.define_vars( + "f", arg_types=[backend.scalar_typ, backend.scalar_typ], ret="scalar" + ) weight = f(u(), v()) idx = (u(), v()) @@ -239,17 +248,10 @@ def test_reduce_dependent_range_mask_delta_body(monoid, reductor, backend: Backe delta(idx, jnp.where(v() < u(), weight, monoid.identity)), {u: Range(0, N, 1), v: Range(0, N, 1)}, ) - - check_rewrite( - lhs=lhs, - rhs=rhs, - rule=ReduceDependentRangeMask(), - backend=backend, - free_vars=[u, v, f], - ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceDependentRangeMask()) -def test_reduce_matmul(): +def test_reduce_matmul(backend: JaxBackend): key = jax.random.PRNGKey(0) # Define dimensions B, I, J, K = 2, 3, 4, 5 @@ -257,7 +259,7 @@ def test_reduce_matmul(): # Create sample matrices X = random.normal(key, (B, I, J)) Y = random.normal(key, (B, J, K)) - (b, i, j, k) = define_vars("b", "i", "j", "k", typ=jax.Array) + (b, i, j, k) = backend.define_vars("b", "i", "j", "k", ret="scalar") with handler(NormalizeIntp): actual = Sum.reduce( diff --git a/tests/test_ops_monoid.py b/tests/test_ops_monoid.py index c7ee7567c..fcd72f064 100644 --- a/tests/test_ops_monoid.py +++ b/tests/test_ops_monoid.py @@ -1,10 +1,13 @@ +import math import typing +from collections.abc import Iterable import pytest from hypothesis import HealthCheck, given, settings from hypothesis import strategies as st import effectful.handlers.jax.monoid # noqa: F401 +import effectful.handlers.jax.numpy as jnp from effectful.ops.monoid import ( CartesianProduct, Max, @@ -22,29 +25,25 @@ PlusSingle, PlusZero, Product, + ReduceCartesianWeightedStream, ReduceDistributeCartesianProduct, ReduceFactorization, ReduceFusion, ReduceNoStreams, ReduceSplit, + ReduceWeightedStream, Sum, distributes_over, ) -from effectful.ops.semantics import fvsof, handler -from effectful.ops.types import Operation -from tests._monoid_helpers import ( - INT_BACKEND, - JAX_BACKEND, - Backend, - check_rewrite, - define_vars, - syntactic_eq_alpha, -) +from effectful.ops.semantics import coproduct, evaluate, fvsof, handler +from effectful.ops.syntax import deffn +from effectful.ops.types import NotHandled, Operation, Term +from tests._monoid_helpers import Backend, IntBackend, JaxBackend, syntactic_eq_alpha -@pytest.fixture(params=[INT_BACKEND, JAX_BACKEND], ids=["int", "jax"]) +@pytest.fixture(params=[IntBackend, JaxBackend], ids=["int", "jax"]) def backend(request) -> Backend: - return request.param + return request.param() ALL_MONOIDS = [ @@ -90,10 +89,10 @@ def backend(request) -> Backend: deadline=None, suppress_health_check=[HealthCheck.function_scoped_fixture], ) -def test_associativity(monoid, backend, data): - a = data.draw(backend.scalar_strategy) - b = data.draw(backend.scalar_strategy) - c = data.draw(backend.scalar_strategy) +def test_associativity(monoid, backend: Backend, data): + a = data.draw(backend.strategy(ret="scalar"))() + b = data.draw(backend.strategy(ret="scalar"))() + c = data.draw(backend.strategy(ret="scalar"))() with handler(NormalizeIntp): left = monoid.plus(monoid.plus(a, b), c) right = monoid.plus(a, monoid.plus(b, c)) @@ -107,8 +106,8 @@ def test_associativity(monoid, backend, data): deadline=None, suppress_health_check=[HealthCheck.function_scoped_fixture], ) -def test_identity(monoid, backend, data): - a = data.draw(backend.scalar_strategy) +def test_identity(monoid, backend: Backend, data): + a = data.draw(backend.strategy(ret="scalar"))() with handler(NormalizeIntp): assert backend.eq(monoid.plus(monoid.identity, a), a) assert backend.eq(monoid.plus(a, monoid.identity), a) @@ -121,9 +120,9 @@ def test_identity(monoid, backend, data): deadline=None, suppress_health_check=[HealthCheck.function_scoped_fixture], ) -def test_commutativity(monoid, backend, data): - a = data.draw(backend.scalar_strategy) - b = data.draw(backend.scalar_strategy) +def test_commutativity(monoid, backend: Backend, data): + a = data.draw(backend.strategy(ret="scalar"))() + b = data.draw(backend.strategy(ret="scalar"))() with handler(NormalizeIntp): assert backend.eq(monoid.plus(a, b), monoid.plus(b, a)) @@ -135,8 +134,8 @@ def test_commutativity(monoid, backend, data): deadline=None, suppress_health_check=[HealthCheck.function_scoped_fixture], ) -def test_idempotence(monoid, backend, data): - a = data.draw(backend.scalar_strategy) +def test_idempotence(monoid, backend: Backend, data): + a = data.draw(backend.strategy(ret="scalar"))() with handler(NormalizeIntp): assert backend.eq(monoid.plus(a, a), a) @@ -148,102 +147,86 @@ def test_idempotence(monoid, backend, data): deadline=None, suppress_health_check=[HealthCheck.function_scoped_fixture], ) -def test_zero_absorbs(monoid, backend, data): - a = data.draw(backend.scalar_strategy) +def test_zero_absorbs(monoid, backend: Backend, data): + a = data.draw(backend.strategy(ret="scalar"))() with handler(NormalizeIntp): assert backend.eq(monoid.plus(monoid.zero, a), monoid.zero) assert backend.eq(monoid.plus(a, monoid.zero), monoid.zero) @pytest.mark.parametrize("monoid", ALL_MONOIDS) -def test_plus_empty(monoid, backend): - check_rewrite( - lhs=monoid.plus(), rhs=monoid.identity, rule=PlusEmpty(), backend=backend - ) +def test_plus_empty(monoid, backend: Backend): + backend.check_rewrite(lhs=monoid.plus(), rhs=monoid.identity, rule=PlusEmpty()) @pytest.mark.parametrize("monoid", ALL_MONOIDS) -def test_plus_single(monoid, backend): - x = define_vars("x", typ=backend.scalar_typ) - check_rewrite( - lhs=monoid.plus(x()), rhs=x(), rule=PlusSingle(), backend=backend, free_vars=[x] - ) +def test_plus_single(monoid, backend: Backend): + x = backend.define_vars("x", ret="scalar") + backend.check_rewrite(lhs=monoid.plus(x()), rhs=x(), rule=PlusSingle()) @pytest.mark.parametrize("monoid", ALL_MONOIDS) -def test_plus_identity_right(monoid, backend): - x = define_vars("x", typ=backend.scalar_typ) +def test_plus_identity_right(monoid, backend: Backend): + x = backend.define_vars("x", ret="scalar") lhs = monoid.plus(x(), monoid.identity) rhs = monoid.plus(x()) - check_rewrite(lhs=lhs, rhs=rhs, rule=PlusIdentity(), backend=backend, free_vars=[x]) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=PlusIdentity()) @pytest.mark.parametrize("monoid", ALL_MONOIDS) -def test_plus_identity_left(monoid, backend): - x = define_vars("x", typ=backend.scalar_typ) +def test_plus_identity_left(monoid, backend: Backend): + x = backend.define_vars("x", ret="scalar") lhs = monoid.plus(monoid.identity, x()) rhs = monoid.plus(x()) - check_rewrite(lhs=lhs, rhs=rhs, rule=PlusIdentity(), backend=backend, free_vars=[x]) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=PlusIdentity()) @pytest.mark.parametrize("monoid", ALL_MONOIDS) -def test_plus_assoc_right(monoid, backend): - x, y, z = define_vars("x", "y", "z", typ=backend.scalar_typ) - check_rewrite( +def test_plus_assoc_right(monoid, backend: Backend): + x, y, z = backend.define_vars("x", "y", "z", ret="scalar") + backend.check_rewrite( lhs=monoid.plus(x(), monoid.plus(y(), z())), rhs=monoid.plus(x(), y(), z()), rule=PlusAssoc(), - backend=backend, - free_vars=[x, y, z], ) @pytest.mark.parametrize("monoid", ALL_MONOIDS) -def test_plus_assoc_left(monoid, backend): - x, y, z = define_vars("x", "y", "z", typ=backend.scalar_typ) - check_rewrite( +def test_plus_assoc_left(monoid, backend: Backend): + x, y, z = backend.define_vars("x", "y", "z", ret="scalar") + backend.check_rewrite( lhs=monoid.plus(monoid.plus(x(), y()), z()), rhs=monoid.plus(x(), y(), z()), rule=PlusAssoc(), - backend=backend, - free_vars=[x, y, z], ) @pytest.mark.parametrize("monoid", ALL_MONOIDS) -def test_plus_sequence(monoid, backend): - a, b, c, d = define_vars("a", "b", "c", "d", typ=backend.scalar_typ) - check_rewrite( +def test_plus_sequence(monoid, backend: Backend): + a, b, c, d = backend.define_vars("a", "b", "c", "d", ret="scalar") + backend.check_rewrite( lhs=monoid.plus((a(), b()), (c(), d())), rhs=(monoid.plus(a(), c()), monoid.plus(b(), d())), rule=MonoidOverSequence(), - backend=backend, - free_vars=[a, b, c, d], ) @pytest.mark.parametrize("monoid", ALL_MONOIDS) -def test_plus_mapping(monoid, backend): - a, b, c, d = define_vars("a", "b", "c", "d", typ=backend.scalar_typ) +def test_plus_mapping(monoid, backend: Backend): + a, b, c, d = backend.define_vars("a", "b", "c", "d", ret="scalar") lhs = monoid.plus({0: a(), 1: b()}, {0: c(), 2: d()}) rhs = {0: monoid.plus(a(), c()), 1: monoid.plus(b()), 2: monoid.plus(d())} - check_rewrite( - lhs=lhs, - rhs=rhs, - rule=MonoidOverMapping(), - backend=backend, - free_vars=[a, b, c, d], - ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=MonoidOverMapping()) -def test_plus_distributes(backend): - a, b, c, d = define_vars("a", "b", "c", "d", typ=backend.scalar_typ) +def test_plus_distributes(backend: Backend): + a, b, c, d = backend.define_vars("a", "b", "c", "d", ret="scalar") lhs = Product.plus(Sum.plus(a(), b()), Sum.plus(c(), d())) rhs = Product.plus( Sum.plus( @@ -253,13 +236,11 @@ def test_plus_distributes(backend): Product.plus(b(), d()), ) ) - check_rewrite( - lhs=lhs, rhs=rhs, rule=PlusDistr(), backend=backend, free_vars=[a, b, c, d] - ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=PlusDistr()) -def test_plus_distributes_constant(backend): - a, b, c, d = define_vars("a", "b", "c", "d", typ=backend.scalar_typ) +def test_plus_distributes_constant(backend: Backend): + a, b, c, d = backend.define_vars("a", "b", "c", "d", ret="scalar") lhs = Product.plus(Sum.plus(a(), b()), Sum.plus(c(), d()), 5) rhs = Product.plus( 5, @@ -270,13 +251,11 @@ def test_plus_distributes_constant(backend): Product.plus(b(), d()), ), ) - check_rewrite( - lhs=lhs, rhs=rhs, rule=PlusDistr(), backend=backend, free_vars=[a, b, c, d] - ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=PlusDistr()) -def test_plus_distributes_multiple(backend): - a, b, c, d = define_vars("a", "b", "c", "d", typ=backend.scalar_typ) +def test_plus_distributes_multiple(backend: Backend): + a, b, c, d = backend.define_vars("a", "b", "c", "d", ret="scalar") lhs = Sum.plus( Min.plus(a(), b()), Min.plus(c(), d()), @@ -297,238 +276,195 @@ def test_plus_distributes_multiple(backend): Sum.plus(b(), d()), ), ) - check_rewrite( - lhs=lhs, rhs=rhs, rule=PlusDistr(), backend=backend, free_vars=[a, b, c, d] - ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=PlusDistr()) @pytest.mark.parametrize("monoid", IDEMPOTENT) -def test_plus_idempotent_consecutive(monoid, backend): +def test_plus_idempotent_consecutive(monoid, backend: Backend): """``a, a, b → a, b`` — only consecutive duplicates collapse.""" - a, b = define_vars("a", "b", typ=backend.scalar_typ) + a, b = backend.define_vars("a", "b", ret="scalar") lhs = monoid.plus(a(), a(), b()) - return check_rewrite( - lhs=lhs, - rhs=monoid.plus(a(), b()), - rule=PlusConsecutiveDups(), - backend=backend, - free_vars=[a, b], + return backend.check_rewrite( + lhs=lhs, rhs=monoid.plus(a(), b()), rule=PlusConsecutiveDups() ) @pytest.mark.parametrize("monoid", IDEMPOTENT) -def test_plus_idempotent_non_consecutive(monoid, backend): +def test_plus_idempotent_non_consecutive(monoid, backend: Backend): """``a, b, a`` — Semilattice (Min/Max) collapses via commutative PlusDups.""" - a, b = define_vars("a", "b", typ=backend.scalar_typ) + a, b = backend.define_vars("a", "b", ret="scalar") lhs = monoid.plus(a(), b(), a()) rhs = monoid.plus(a(), b()) - check_rewrite(lhs=lhs, rhs=rhs, rule=PlusDups(), backend=backend, free_vars=[a, b]) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=PlusDups()) @pytest.mark.parametrize("monoid", [Min, Max]) -def test_plus_commutative_idempotent_long(monoid, backend): +def test_plus_commutative_idempotent_long(monoid, backend: Backend): """Long alternation collapses via commutative dedup (Min/Max only).""" - a, b = define_vars("a", "b", typ=backend.scalar_typ) + a, b = backend.define_vars("a", "b", ret="scalar") lhs = monoid.plus(a(), b(), a(), b(), b(), a(), a()) rhs = monoid.plus(a(), b()) - check_rewrite(lhs=lhs, rhs=rhs, rule=PlusDups(), backend=backend, free_vars=[a, b]) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=PlusDups()) @pytest.mark.parametrize("monoid", WITH_ZERO) -def test_plus_zero(monoid, backend): - a = define_vars("a", typ=backend.scalar_typ) +def test_plus_zero(monoid, backend: Backend): + a = backend.define_vars("a", ret="scalar") lhs_right = monoid.plus(a(), monoid.zero) lhs_left = monoid.plus(monoid.zero, a()) rhs = monoid.zero - check_rewrite( - lhs=lhs_right, rhs=rhs, rule=PlusZero(), backend=backend, free_vars=[a] - ) - check_rewrite( - lhs=lhs_left, rhs=rhs, rule=PlusZero(), backend=backend, free_vars=[a] - ) + backend.check_rewrite(lhs=lhs_right, rhs=rhs, rule=PlusZero()) + backend.check_rewrite(lhs=lhs_left, rhs=rhs, rule=PlusZero()) @pytest.mark.parametrize("monoid", ALL_MONOIDS) -def test_partial_1(monoid, backend): - x, y = define_vars("x", "y", typ=backend.scalar_typ) +def test_partial_1(monoid, backend: Backend): + x = backend.define_vars("x", ret="scalar") lhs = monoid.reduce(x(), {x: []}) rhs = monoid.identity - check_rewrite(lhs=lhs, rhs=rhs, rule={}, backend=backend, free_vars=[x, y]) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule={}) @pytest.mark.parametrize("monoid", ALL_MONOIDS) -def test_partial_2(monoid, backend): - x, y = define_vars("x", "y", typ=backend.scalar_typ) - Y = define_vars("Y", typ=backend.stream_typ) +def test_partial_2(monoid, backend: Backend): + x, y = backend.define_vars("x", "y", ret="scalar") + Y = backend.define_vars("Y", ret="stream") lhs = monoid.reduce(x(), {y: Y(), x: []}) rhs = monoid.identity - - check_rewrite(lhs=lhs, rhs=rhs, rule={}, backend=backend, free_vars=[x, y, Y]) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule={}) @pytest.mark.parametrize("monoid", ALL_MONOIDS) -def test_partial_3(monoid, backend): - x, y, a, b = define_vars("x", "y", "a", "b", typ=backend.scalar_typ) - Y = define_vars("Y", typ=backend.stream_typ) +def test_partial_3(monoid, backend: Backend): + x, y, a, b = backend.define_vars("x", "y", "a", "b", ret="scalar") + Y = backend.define_vars("Y", ret="stream") lhs = monoid.reduce(x(), {y: Y(), x: [a(), b()]}) rhs = monoid.plus(monoid.reduce(a(), {y: Y()}), monoid.reduce(b(), {y: Y()})) - - check_rewrite(lhs=lhs, rhs=rhs, rule={}, backend=backend, free_vars=[x, y, a, b, Y]) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule={}) @pytest.mark.parametrize("monoid", ALL_MONOIDS) -def test_partial_4(monoid, backend): - x, y, a, b = define_vars("x", "y", "a", "b", typ=backend.scalar_typ) - f = backend.fresh_op("f", n_args=1, ret="stream") +def test_partial_4(monoid, backend: Backend): + x, y, a, b = backend.define_vars("x", "y", "a", "b", ret="scalar") + f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="stream") lhs = monoid.reduce(x(), {y: f(x()), x: [a(), b()]}) rhs = monoid.plus(monoid.reduce(a(), {y: f(a())}), monoid.reduce(b(), {y: f(b())})) - - check_rewrite(lhs=lhs, rhs=rhs, rule={}, backend=backend, free_vars=[x, y, a, b, f]) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule={}) @pytest.mark.parametrize("monoid", ALL_MONOIDS) -def test_reduce_body_sequence(monoid, backend): - x = Operation.define(backend.scalar_typ, name="x") - X = Operation.define(backend.stream_typ, name="X") - f = backend.fresh_op("f", n_args=1, ret="scalar") - g = Operation.define(f, name="g") +def test_reduce_body_sequence(monoid, backend: Backend): + x = backend.define_vars("x", ret="scalar") + X = backend.define_vars("X", ret="stream") + f, g = backend.define_vars("f", "g", arg_types=(backend.scalar_typ,), ret="scalar") lhs = monoid.reduce((f(x()), g(x())), {x: X()}) rhs = (monoid.reduce(f(x()), {x: X()}), monoid.reduce(g(x()), {x: X()})) - - check_rewrite( - lhs=lhs, - rhs=rhs, - rule=MonoidOverSequence(), - backend=backend, - free_vars=[X, f, g], - ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=MonoidOverSequence()) @pytest.mark.parametrize("monoid", ALL_MONOIDS) -def test_reduce_body_sequence_2(monoid, backend): - x, y = define_vars("x", "y", typ=backend.scalar_typ) - X, Y = define_vars("X", "Y", typ=backend.stream_typ) - f = backend.fresh_op("f", n_args=1, ret="scalar") - g = Operation.define(f, name="g") +def test_reduce_body_sequence_2(monoid, backend: Backend): + x, y = backend.define_vars("x", "y", ret="scalar") + X, Y = backend.define_vars("X", "Y", ret="stream") + f, g = backend.define_vars("f", "g", arg_types=(backend.scalar_typ,), ret="scalar") lhs = monoid.reduce((f(x()), g(y())), {x: X(), y: Y()}) rhs = ( monoid.reduce(f(x()), {x: X(), y: Y()}), monoid.reduce(g(y()), {x: X(), y: Y()}), ) - - check_rewrite( - lhs=lhs, - rhs=rhs, - rule=MonoidOverSequence(), - backend=backend, - free_vars=[X, Y, f, g], - ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=MonoidOverSequence()) @pytest.mark.parametrize("monoid", ALL_MONOIDS) -def test_reduce_body_mapping(monoid, backend): - x = Operation.define(backend.scalar_typ, name="x") - X = Operation.define(backend.stream_typ, name="X") - f = backend.fresh_op("f", n_args=1, ret="scalar") - g = Operation.define(f, name="g") +def test_reduce_body_mapping(monoid, backend: Backend): + x = backend.define_vars("x", ret="scalar") + X = backend.define_vars("X", ret="stream") + f, g = backend.define_vars("f", "g", arg_types=(backend.scalar_typ,), ret="scalar") lhs = monoid.reduce({0: f(x()), 1: g(x())}, {x: X()}) rhs = { 0: monoid.reduce(f(x()), {x: X()}), 1: monoid.reduce(g(x()), {x: X()}), } - check_rewrite( - lhs=lhs, - rhs=rhs, - rule=MonoidOverMapping(), - backend=backend, - free_vars=[X, f, g], - ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=MonoidOverMapping()) @pytest.mark.parametrize("monoid", ALL_MONOIDS) -def test_reduce_no_streams(monoid, backend): - a = define_vars("a", typ=backend.scalar_typ) +def test_reduce_no_streams(monoid, backend: Backend): + a = backend.define_vars("a", ret="scalar") + lhs = monoid.reduce(a(), {}) rhs = monoid.identity - - check_rewrite( - lhs=lhs, rhs=rhs, rule=ReduceNoStreams(), backend=backend, free_vars=[a] - ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceNoStreams()) @pytest.mark.parametrize("monoid", ALL_MONOIDS) -def test_reduce_reduce(monoid, backend): - a, b = define_vars("a", "b", typ=backend.scalar_typ) - A, B = define_vars("A", "B", typ=backend.stream_typ) - f = backend.fresh_op("f", n_args=2, ret="scalar") +def test_reduce_reduce(monoid, backend: Backend): + a, b = backend.define_vars("a", "b", ret="scalar") + A, B = backend.define_vars("A", "B", ret="stream") + f = backend.define_vars( + "f", arg_types=(backend.scalar_typ, backend.scalar_typ), ret="scalar" + ) lhs = monoid.reduce(monoid.reduce(f(a(), b()), {a: A()}), {b: B()}) rhs = monoid.reduce(f(a(), b()), {a: A(), b: B()}) - - check_rewrite( - lhs=lhs, rhs=rhs, rule=ReduceFusion(), backend=backend, free_vars=[A, B, f] - ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceFusion()) @pytest.mark.parametrize("monoid", COMMUTATIVE) -def test_reduce_plus(monoid, backend): - a, b = define_vars("a", "b", typ=backend.scalar_typ) - A, B = define_vars("A", "B", typ=backend.stream_typ) +def test_reduce_plus(monoid, backend: Backend): + a, b = backend.define_vars("a", "b", ret="scalar") + A, B = backend.define_vars("A", "B", ret="stream") + lhs = monoid.reduce(monoid.plus(a(), b()), {a: A(), b: B()}) rhs = monoid.plus( monoid.reduce(a(), {a: A(), b: B()}), monoid.reduce(b(), {a: A(), b: B()}), ) - check_rewrite( - lhs=lhs, rhs=rhs, rule=ReduceSplit(), backend=backend, free_vars=[A, B] - ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceSplit()) -def test_reduce_independent_1(backend): - a, b = define_vars("a", "b", typ=backend.scalar_typ) - A, B = define_vars("A", "B", typ=backend.stream_typ) +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") + lhs = Sum.reduce(Product.plus(a(), b()), {a: A(), b: B()}) rhs = Product.plus( Sum.reduce(Product.plus(a()), {a: A()}), Sum.reduce(Product.plus(b()), {b: B()}) ) - check_rewrite( - lhs=lhs, rhs=rhs, rule=ReduceFactorization(), backend=backend, free_vars=[A, B] - ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceFactorization()) -def test_reduce_independent_2(backend): - a, b, c = define_vars("a", "b", "c", typ=backend.scalar_typ) - A, B, C = define_vars("A", "B", "C", typ=backend.stream_typ) - f = backend.fresh_op("f", n_args=2, ret="scalar") +def test_reduce_independent_2(backend: Backend): + a, b, c = backend.define_vars("a", "b", "c", ret="scalar") + A, B, C = backend.define_vars("A", "B", "C", ret="stream") + f = backend.define_vars( + "f", arg_types=(backend.scalar_typ, backend.scalar_typ), ret="scalar" + ) lhs = Sum.reduce(Product.plus(a(), b(), f(b(), c())), {a: A(), b: B(), c: C()}) rhs = Product.plus( Sum.reduce(Product.plus(a()), {a: A()}), Sum.reduce(Product.plus(b(), f(b(), c())), {b: B(), c: C()}), ) - check_rewrite( - lhs=lhs, - rhs=rhs, - rule=ReduceFactorization(), - backend=backend, - free_vars=[A, B, C, f], - ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceFactorization()) -def test_reduce_independent_3_negative(backend): +def test_reduce_independent_3_negative(backend: Backend): """Stream `b` depends on `a` (b: g(a())), so the proposed factorization is unsound — the normalizer must NOT apply it.""" - a, b, c = define_vars("a", "b", "c", typ=backend.scalar_typ) - A, C = define_vars("A", "C", typ=backend.stream_typ) - f = backend.fresh_op("f", n_args=2, ret="scalar") - g = backend.fresh_op("g", n_args=1, ret="stream") + a, b, c = backend.define_vars("a", "b", "c", ret="scalar") + A, C = backend.define_vars("A", "C", ret="stream") + f = backend.define_vars( + "f", arg_types=(backend.scalar_typ, backend.scalar_typ), ret="scalar" + ) + g = backend.define_vars("g", arg_types=(backend.scalar_typ,), ret="stream") with handler(ReduceFactorization()): # ty:ignore[invalid-argument-type] lhs = Sum.reduce( @@ -542,10 +478,12 @@ def test_reduce_independent_3_negative(backend): assert not syntactic_eq_alpha(lhs, bogus_rhs) -def test_reduce_independent_4(backend): - a, b, c = define_vars("a", "b", "c", typ=backend.scalar_typ) - A, B, C = define_vars("A", "B", "C", typ=backend.stream_typ) - f = backend.fresh_op("f", n_args=2, ret="scalar") +def test_reduce_independent_4(backend: Backend): + a, b, c = backend.define_vars("a", "b", "c", ret="scalar") + A, B, C = backend.define_vars("A", "B", "C", ret="stream") + f = backend.define_vars( + "f", arg_types=(backend.scalar_typ, backend.scalar_typ), ret="scalar" + ) lhs = Sum.reduce(Product.plus(a(), b(), f(b(), c()), 7), {a: A(), b: B(), c: C()}) rhs = Product.plus( @@ -553,39 +491,44 @@ def test_reduce_independent_4(backend): Sum.reduce(Product.plus(a()), {a: A()}), Sum.reduce(Product.plus(b(), f(b(), c())), {b: B(), c: C()}), ) - check_rewrite( - lhs=lhs, - rhs=rhs, - rule=ReduceFactorization(), - backend=backend, - free_vars=[A, B, C, f], - ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceFactorization()) + + +def test_reduce_cartesian_3(): + backend = JaxBackend() + i = backend.define_vars("i", ret="scalar") + + with handler(NormalizeIntp): + value = CartesianProduct.reduce(jnp.zeros(2), {i: jnp.arange(3)}) + assert value.shape == (2**3, 3) + + with handler(NormalizeIntp): + value = CartesianProduct.reduce(jnp.zeros(2), {i: jnp.arange(1)}) + assert value.shape == (2**1, 1) + + with handler(NormalizeIntp): + value = CartesianProduct.reduce(jnp.zeros(1), {i: jnp.arange(3)}) + assert value.shape == (1**3, 3) @pytest.mark.parametrize("outer,inner", MONOID_PAIRS) -def test_reduce_lifted_1(outer, inner, backend): - a, i = define_vars("a", "i", typ=backend.scalar_typ) - A, N, A_domain = define_vars("A", "N", "A_domain", typ=backend.stream_typ) - f = backend.fresh_op("f", n_args=1, ret="scalar") +def test_reduce_lifted_1(outer, inner, backend: Backend): + a, i = backend.define_vars("a", "i", ret="scalar") + A, N, A_domain = backend.define_vars("A", "N", "A_domain", ret="stream") + f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="scalar") - term1 = outer.reduce( + lhs = outer.reduce( inner.reduce(f(a()), {a: A()}), {A: CartesianProduct.reduce(A_domain(), {i: N()})}, ) - term2 = inner.reduce(outer.reduce(inner.plus(f(a())), {a: A_domain()}), {i: N()}) - - check_rewrite( - lhs=term1, - rhs=term2, - rule=ReduceDistributeCartesianProduct(), - backend=backend, - free_vars=[N, A_domain, f], - ) + rhs = inner.reduce(outer.reduce(inner.plus(f(a())), {a: A_domain()}), {i: N()}) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceDistributeCartesianProduct()) def test_reduce_cartesian_1(): - a, i = define_vars("a", "i", typ=int) - A = define_vars("A", typ=tuple[int]) + backend = IntBackend() + a, i = backend.define_vars("a", "i", ret="scalar") + A = backend.define_vars("A", ret="stream") with handler(NormalizeIntp): term1 = Sum.reduce( @@ -597,8 +540,9 @@ def test_reduce_cartesian_1(): def test_reduce_cartesian_2(): - a, i = define_vars("a", "i", typ=int) - A = define_vars("A", typ=tuple[int]) + backend = IntBackend() + a, i = backend.define_vars("a", "i", ret="scalar") + A = backend.define_vars("A", ret="stream") with handler(NormalizeIntp): term1 = Sum.reduce( @@ -610,46 +554,41 @@ def test_reduce_cartesian_2(): @pytest.mark.parametrize("outer,inner", MONOID_PAIRS) -def test_reduce_lifted_multi_index(outer, inner, backend): - a, i, j = define_vars("a", "i", "j", typ=backend.scalar_typ) - A, N, M, A_domain = define_vars("A", "N", "M", "A_domain", typ=backend.stream_typ) - f = backend.fresh_op("f", n_args=1, ret="scalar") +def test_reduce_lifted_multi_index(outer, inner, backend: Backend): + a, i, j = backend.define_vars("a", "i", "j", ret="scalar") + A, N, M, A_domain = backend.define_vars("A", "N", "M", "A_domain", ret="stream") + f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="scalar") - term1 = outer.reduce( + lhs = outer.reduce( inner.reduce(f(a()), {a: A()}), {A: CartesianProduct.reduce(A_domain(), {i: N(), j: M()})}, ) - term2 = inner.reduce( - outer.reduce(inner.plus(f(a())), {a: A_domain()}), - {i: N(), j: M()}, - ) - check_rewrite( - lhs=term1, - rhs=term2, - rule=ReduceDistributeCartesianProduct(), - backend=backend, - free_vars=[N, M, A_domain, f], + rhs = inner.reduce( + outer.reduce(inner.plus(f(a())), {a: A_domain()}), {i: N(), j: M()} ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceDistributeCartesianProduct()) @pytest.mark.parametrize("outer,inner", MONOID_PAIRS) -def test_reduce_lifted_2(outer, inner, backend): +def test_reduce_lifted_2(outer, inner, backend: Backend): """The worked example on page 396 of 'Lifted Variable Elimination: Decoupling the Operators from the Constraint Language'. """ - a, i, s, t = define_vars("a", "i", "s", "t", typ=backend.scalar_typ) - A, N, T = define_vars("A", "N", "T", typ=backend.stream_typ) - A_domain = backend.fresh_op("A_domain", n_args=1, ret="stream") - f1 = backend.fresh_op("f1", n_args=2, ret="scalar") - f2 = backend.fresh_op("f2", n_args=2, ret="scalar") + a, i, s, t = backend.define_vars("a", "i", "s", "t", ret="scalar") + A, N, T = backend.define_vars("A", "N", "T", ret="stream") + A_domain = backend.define_vars( + "A_domain", arg_types=(backend.scalar_typ,), ret="stream" + ) + f1, f2 = backend.define_vars( + "f1", "f2", arg_types=(backend.scalar_typ, backend.scalar_typ), ret="scalar" + ) - term1 = outer.reduce( + lhs = outer.reduce( inner.reduce(inner.plus(f1(a(), s()), f2(t(), a())), {a: A()}), {A: CartesianProduct.reduce(A_domain(i()), {i: N()}), t: T()}, ) - - term2 = outer.reduce( + rhs = outer.reduce( inner.reduce( outer.reduce( inner.plus(inner.plus(f1(a(), s()), f2(t(), a()))), {a: A_domain(i())} @@ -658,11 +597,143 @@ def test_reduce_lifted_2(outer, inner, backend): ), {t: T()}, ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceDistributeCartesianProduct()) + + +# --------------------------------------------------------------------------- +# Weighted streams +# --------------------------------------------------------------------------- + + +def test_reduce_single_weighted_stream(backend: Backend): + """Single weighted stream desugars: + Sum.reduce(body, {a: WS(A, w, Product)}) + = Sum.reduce(Product.plus(w(a), body), {a: A}) + """ + a = backend.define_vars("a", ret="scalar") + A = backend.define_vars("A", ret="stream") + body, w = backend.define_vars( + "body", "w", arg_types=(backend.scalar_typ,), ret="scalar" + ) + + lhs = Sum.reduce(body(a()), {a: Product.weighted(A(), w)}) + rhs = Sum.reduce(Product.plus(w(a()), body(a())), {a: A()}) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceWeightedStream()) + + +def test_reduce_weighted_factorization(backend: Backend): + """Two independent weighted streams under Sum with Product weights factor: + Sum.reduce(f(a)*g(b), {a: Product.weighted(A, a, w_a), b: Product.weighted(B, b, w_b)}) + = (Sum.reduce(w_a(a)*f(a), {a: A})) * (Sum.reduce(w_b(b)*g(b), {b: B})) + + Exercises chaining of ``ReduceWeightedStream`` with ``ReduceFactorization`` + inside ``NormalizeIntp``. + """ + a, b = backend.define_vars("a", "b", ret="scalar") + A, B = backend.define_vars("A", "B", ret="stream") + f, g, w_a, w_b = backend.define_vars( + "f", "g", "w_a", "w_b", arg_types=(backend.scalar_typ,), ret="scalar" + ) + + lhs = Sum.reduce( + Product.plus(f(a()), g(b())), + {a: Product.weighted(A(), w_a), b: Product.weighted(B(), w_b)}, + ) + rhs = Product.plus( + Sum.reduce(Product.plus(w_a(a()), Product.plus(f(a()))), {a: A()}), + Sum.reduce(Product.plus(w_b(b()), Product.plus(g(b()))), {b: B()}), + ) + backend.check_rewrite( + lhs=lhs, rhs=rhs, rule=coproduct(ReduceWeightedStream(), ReduceFactorization()) + ) + + +def test_reduce_cartesian_weighted_stream(backend: Backend): + """``CartesianProduct.reduce`` over a ``WeightedStream`` body whose weight + is independent of the plate var rewrites to a single joint + ``WeightedStream``: + + CartesianProduct.reduce(M.weighted(s, e, w(e)), {p: P}) + = M.weighted(CartesianProduct.reduce(s, {p: P}), row, M.reduce(w(e), {e: row()})) + """ + p, e_var = backend.define_vars("p", "e_var", ret="scalar") + S, P = backend.define_vars("S", "P", ret="stream") + w = backend.define_vars("w", arg_types=(backend.scalar_typ,), ret="scalar") + + lhs = CartesianProduct.reduce(Product.weighted(S(), w), {p: P()}) + row_var = Operation.define(Iterable[backend.scalar_typ], name="row") # type: ignore[name-defined] + rhs = Product.weighted( + CartesianProduct.reduce(S(), {p: P()}), + deffn(Product.reduce(w(e_var()), {e_var: row_var()}), row_var), + ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceCartesianWeightedStream()) + + +def test_lift_weighted_cartesian(backend: Backend): + """Compose ``ReduceCartesianWeightedStream`` + ``ReduceWeightedStream`` + + ``ReduceDistributeCartesianProduct`` on a Sum-of-Product-of-weighted shape: + + Sum.reduce( + Product.reduce(body(a()), {a: A()}), + {A: CartesianProduct.reduce(Product.weighted(S, e, w(e)), {p: P})}, + ) + + The inner ``weighted`` becomes a joint ``weighted`` (rule 1), lifts its + per-element weight into the outer Sum body (rule 2), and the lifted form + matches the inversion pattern (rule 3), yielding:: + + Product.reduce( + Sum.reduce(Product.plus(w(a()), body(a())), {a: S}), + {p: P}, + ) + """ + a, p = backend.define_vars("a", "p", ret="scalar") + A, S, P = backend.define_vars("A", "S", "P", ret="stream") + body, w = backend.define_vars( + "body", "w", arg_types=(backend.scalar_typ,), ret="scalar" + ) - check_rewrite( - lhs=term1, - rhs=term2, - rule=ReduceDistributeCartesianProduct(), - backend=backend, - free_vars=[a, i, s, t, A, N, T, A_domain, f1, f2], + lhs = Sum.reduce( + Product.reduce(body(a()), {a: A()}), + {A: CartesianProduct.reduce(Product.weighted(S(), w), {p: P()})}, + ) + rhs = Product.reduce( + Sum.reduce(Product.plus(w(a()), body(a())), {a: S()}), {p: P()} ) + backend.check_rewrite( + lhs=lhs, + rhs=rhs, + rule=coproduct( + coproduct(ReduceWeightedStream(), ReduceCartesianWeightedStream()), + ReduceDistributeCartesianProduct(), + ), + ) + + +def test_weighted_expectation_demo(): + """Demo: compute E[f(X)] = Σ_x w(x)·f(x) via a weighted reduce. + + X ranges over [1, 2, 3, 4] with weights w(x) = x/10 (a valid distribution + since the weights sum to 1) and f(x) = x*x. Expected value: + 0.1·1 + 0.2·4 + 0.3·9 + 0.4·16 = 10.0 + """ + weights = {1: 0.1, 2: 0.2, 3: 0.3, 4: 0.4} + + def _w(v: int) -> float: + if isinstance(v, Term): + raise NotHandled + return weights[v] + + def _f(v: int) -> float: + if isinstance(v, Term): + raise NotHandled + return float(v * v) + + a = Operation.define(int, name="a") + w = Operation.define(_w, name="w") + f = Operation.define(_f, name="f") + + with handler(NormalizeIntp): + result = evaluate(Sum.reduce(f(a()), {a: Product.weighted([1, 2, 3, 4], w)})) + + assert math.isclose(result, 10.0) From fb9ec75edea3dcb0e46dde3c6301239ff148be8c Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Mon, 8 Jun 2026 12:08:22 -0400 Subject: [PATCH 07/16] Replace factorization rules with a push-based rule (#672) * more agressive factorization that hoists shared streams * reduce nesting * comment * replace with simpler push-based rule * format * drop unused disjoint set * remove unused * push multiple streams instead of one at a time --- effectful/internals/disjoint_set.py | 99 ------------- effectful/ops/monoid.py | 205 +++++++++++++++------------ tests/_monoid_helpers.py | 10 ++ tests/test_internals_disjoint_set.py | 124 ---------------- tests/test_ops_monoid.py | 77 +++++++++- 5 files changed, 198 insertions(+), 317 deletions(-) delete mode 100644 effectful/internals/disjoint_set.py delete mode 100644 tests/test_internals_disjoint_set.py diff --git a/effectful/internals/disjoint_set.py b/effectful/internals/disjoint_set.py deleted file mode 100644 index 73b5c5c52..000000000 --- a/effectful/internals/disjoint_set.py +++ /dev/null @@ -1,99 +0,0 @@ -class DisjointSet: - """Disjoint Set Union (Union-Find) data structure. - - Maintains a collection of disjoint sets over the integers 0..n-1, - supporting near-constant-time union and find operations via - path compression and union by rank. - - The amortized time complexity per operation is O(α(n)), where α - is the inverse Ackermann function (effectively constant for any - practical n). - - Example: - >>> dsu = DisjointSet(5) - >>> dsu.union(0, 1) - True - >>> dsu.union(1, 2) - True - >>> dsu.find(0) == dsu.find(2) - True - >>> dsu.find(0) == dsu.find(3) - False - """ - - def __init__(self, n): - """Initialize n singleton sets: {0}, {1}, ..., {n-1}. - - Args: - n: The number of elements. Elements are labeled 0..n-1. - """ - self.parent = list(range(n)) - self.rank = [0] * n - - def _validate(self, x): - if x < 0 or x >= len(self.parent): - raise IndexError(f"Element {x} out of bounds") - - def find(self, x): - """Return the representative (root) of the set containing x. - - Two elements belong to the same set if and only if they have - the same representative. Applies path compression: every node - traversed is re-parented directly to its grandparent, flattening - the tree to speed up future queries. - - Args: - x: The element to look up. - - Returns: - The root element of x's set. - """ - self._validate(x) - while self.parent[x] != x: - self.parent[x] = self.parent[self.parent[x]] # path compression - x = self.parent[x] - return x - - def union(self, *elements): - """Merge the sets containing all given elements into one. - - Accepts any number of elements and unions them all together. - Uses union by rank: shallower trees are attached under the root - of the deeper one, keeping the combined tree shallow. - - Args: - *elements: Two or more elements to merge into a single set. - Calling with 0 or 1 elements is a no-op and returns False. - - Returns: - True if any merging occurred (i.e., at least two of the - elements were in different sets); False if all elements - were already in the same set or fewer than 2 were given. - """ - if len(elements) < 2: - return False - - merged = False - first = elements[0] - - for y in elements[1:]: - if self._union_pair(first, y): - merged = True - - return merged - - def _union_pair(self, x, y): - rx = self.find(x) - ry = self.find(y) - - if rx == ry: - return False - - if self.rank[rx] < self.rank[ry]: - rx, ry = ry, rx - - self.parent[ry] = rx - if self.rank[rx] == self.rank[ry]: - self.rank[rx] += 1 - - return True diff --git a/effectful/ops/monoid.py b/effectful/ops/monoid.py index 76351fa62..5f342f25b 100644 --- a/effectful/ops/monoid.py +++ b/effectful/ops/monoid.py @@ -9,7 +9,6 @@ from graphlib import TopologicalSorter from typing import Annotated, Any -from effectful.internals.disjoint_set import DisjointSet from effectful.ops.semantics import ( coproduct, evaluate, @@ -62,6 +61,51 @@ def outer_stream(streams: Streams) -> Iterable[tuple[Operation, Stream, Streams] ) +def inner_stream( + streams: dict[Operation, Expr], +) -> Iterable[tuple[dict[Operation, Expr], Operation, Expr]]: + """Returns the streams that can be ordered innermost in the loop nest as + well as the remaining streams in the nest. + + """ + stream_vars = set(streams.keys()) + + no_dependents = set() + succ = defaultdict(set) + for k, v in streams.items(): + preds = fvsof(v) & stream_vars + if preds: + for pred in preds: + succ[pred].add(k) + else: + no_dependents.add(k) + + topo = TopologicalSorter(succ) + topo.prepare() + return ( + ({k: v for (k, v) in streams.items() if k != op}, op, streams[op]) + for op in set(topo.get_ready()) | no_dependents + ) + + +def inner_streams_first(streams: dict[Operation, Expr]) -> Iterable[Operation]: + """Iterable over streams where dependent streams precede their dependencies.""" + stream_vars = set(streams.keys()) + + no_dependents = set() + succ = defaultdict(set) + for k, v in streams.items(): + preds = fvsof(v) & stream_vars + if preds: + for pred in preds: + succ[pred].add(k) + else: + no_dependents.add(k) + + topo = TopologicalSorter(succ) + return topo.static_order() + + class Monoid[W]: """A monoid with ``plus`` and ``reduce`` :class:`Operation` s.""" @@ -392,110 +436,87 @@ def reduce(self, monoid, body, streams): class ReduceFactorization(ObjectInterpretation): - """ - Implements factorization of independent terms. - For example, when having two independent distributions, - we can rewrite their marginalization as: - ∫p(x)⋅q(y)dxdy => ∫p(x)dx ⋅ ∫q(y)dy - - More specifically, in terms of reduces we are performing: - reduce(R, (S₁ × ... × Sₖ) , A₁ * ... * Aₖ) - => reduce(R, S₁, A₁) * ... * reduce(R, Sₖ, Aₖ) - where free(Aᵢ) ∩ free(Aⱼ) ∩ S = ∅ - and free(Aᵢ) ∩ S ⊆ Sᵢ + """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 + v has no dependents among the remaining streams (so it can be innermost) + and F_rest is nonempty (universal variables stay in the outer core). """ @implements(Monoid.reduce) def reduce(self, monoid, body, streams): - if not is_commutative(monoid): - return fwd() - if ( - isinstance(body, Term) + if not ( + is_commutative(monoid) + and isinstance(body, Term) and _is_monoid_plus(body.op) and distributes_over(body.op.__self__, monoid) ): - inner_monoid: Monoid = body.op.__self__ - stream_vars = set(streams.keys()) - factors = [(arg, fvsof(arg)) for arg in body.args] - stream_ids = {v: i for (i, v) in enumerate(stream_vars)} - ds = DisjointSet(len(streams)) - - # streams are in the same partition as their dependencies - for stream_var, stream_id in stream_ids.items(): - stream_body = streams[stream_var] - deps = sorted([stream_ids[v] for v in fvsof(stream_body) & stream_vars]) - ds.union(stream_id, *deps) - - # factors are in the same partition as their dependencies - for _, factor_fvs in factors: - factor_streams = sorted( - [stream_ids[v] for v in (factor_fvs & stream_vars)] - ) - ds.union(*factor_streams) - - placed_streams = set() - new_reduces = [] - for stream_key in streams: - if stream_key in placed_streams: - continue - - partition = ds.find(stream_ids[stream_key]) - partition_streams = { - k: v - for (k, v) in streams.items() - if ds.find(stream_ids[k]) == partition - } - partition_stream_keys = set(partition_streams.keys()) - - partition_factors = [ - t for t in factors if (t[1] & partition_stream_keys) - ] - - assert all( - (t[1] & stream_vars) <= partition_stream_keys - for t in partition_factors - ), "partition contains all streams required by factor" - - partition_term = inner_monoid.plus(*(t[0] for t in partition_factors)) - new_reduces.append((partition_term, partition_streams)) - placed_streams |= partition_stream_keys - - constant_factors = [t for (t, fvs) in factors if not (fvs & stream_vars)] - - if len(new_reduces) > 1: - result = inner_monoid.plus( - *constant_factors, *(monoid.reduce(*args) for args in new_reduces) - ) - return result + return fwd() - return fwd() + inner = body.op.__self__ + stream_keys = set(streams) + factors = [(a, fvsof(a)) for a in body.args] + # candidates: innermost-eligible (no remaining stream depends on v), + # non-universal (some factor doesn't mention v) + support: dict = {} + for v in streams: + if any(v in fvsof(s) for k, s in streams.items() if k is not v): + continue + f_v = frozenset(i for i, (_, fvs) in enumerate(factors) if v in fvs) + if len(f_v) == len(factors): + continue # v is universal: leave it in the outer core + support[v] = f_v + + # eliminate a variable with subset-minimal factor support + # (leaves-first; canonical on hierarchical/laminar supports) + inner_stream = None + inner_factor_ids = None + for v, f_v in support.items(): + if any(u_sup < f_v for u, u_sup in support.items() if u is not v): + continue + inner_stream = v + inner_factor_ids = f_v + break -def inner_stream( - streams: dict[Operation, Expr], -) -> Iterable[tuple[dict[Operation, Expr], Operation, Expr]]: - """Returns the streams that can be ordered innermost in the loop nest as - well as the remaining streams in the nest. + if not inner_stream or not inner_factor_ids: + return fwd() - """ - stream_vars = set(streams.keys()) + inner_factors = [factors[i][0] for i in sorted(inner_factor_ids)] + inner_stream_keys = {inner_stream} + inner_deps = set().union( + *(factors[i][1] for i in f_v), fvsof(streams[v]) & stream_keys + ) - no_dependents = set() - succ = defaultdict(set) - for k, v in streams.items(): - preds = fvsof(v) & stream_vars - if preds: - for pred in preds: - succ[pred].add(k) - else: - no_dependents.add(k) + outer_factors = [a for i, (a, _) in enumerate(factors) if i not in f_v] + outer_stream_keys = stream_keys - inner_stream_keys + outer_factor_deps = set().union( + *(vars for i, (_, vars) in enumerate(factors) if i not in f_v) + ) - topo = TopologicalSorter(succ) - topo.prepare() - return ( - ({k: v for (k, v) in streams.items() if k != op}, op, streams[op]) - for op in set(topo.get_ready()) | no_dependents - ) + # find all streams that are used in the inner factors/streams and are + # not used by the outer factors/streams + # this has to be done iteratively, because moving a stream inward + # reduces the outer dependency set + # ensures that no future factorization application creates a reduce that + # fuses with with the inner reduce + for s in inner_streams_first(streams): + outer_stream_deps = ( + set().union(*(fvsof(streams[k]) for k in outer_stream_keys)) + & stream_keys + ) + outer_deps = outer_factor_deps | outer_stream_deps + if s in inner_deps and s not in outer_deps: + inner_stream_keys |= {s} + inner_deps |= stream_keys & fvsof(streams[s]) + 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) + + rest_streams = {k: s for k, s in streams.items() if k in outer_stream_keys} + new_body = inner.plus(*outer_factors, inner_red) + return monoid.reduce(new_body, rest_streams) if rest_streams else new_body class ReduceDistributeCartesianProduct(ObjectInterpretation): diff --git a/tests/_monoid_helpers.py b/tests/_monoid_helpers.py index f8089bec7..72787558a 100644 --- a/tests/_monoid_helpers.py +++ b/tests/_monoid_helpers.py @@ -322,6 +322,11 @@ def strategy( return st.sampled_from(self._unary_num_fns) case (builtins.int, builtins.int), "scalar": return st.sampled_from(self._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), + ).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) raise NotImplementedError( @@ -389,6 +394,11 @@ def strategy( return st.sampled_from(self._unary_jax_scalar_fns) case (jax.Array, jax.Array), "scalar": return st.sampled_from(self._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), + ).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) diff --git a/tests/test_internals_disjoint_set.py b/tests/test_internals_disjoint_set.py deleted file mode 100644 index 808b8d25d..000000000 --- a/tests/test_internals_disjoint_set.py +++ /dev/null @@ -1,124 +0,0 @@ -import random - -import pytest - -from effectful.internals.disjoint_set import DisjointSet - - -@pytest.fixture -def dsu(): - return DisjointSet(10) - - -def test_initial_state(dsu): - for i in range(10): - assert dsu.find(i) == i - - -def test_simple_union(dsu): - assert dsu.union(1, 2) is True - assert dsu.find(1) == dsu.find(2) - - -def test_union_idempotent(dsu): - dsu.union(1, 2) - assert dsu.union(1, 2) is False - - -def test_union_chain(dsu): - dsu.union(1, 2) - dsu.union(2, 3) - assert dsu.find(1) == dsu.find(3) - - -def test_union_multiple_elements_all_connected(dsu): - dsu.union(1, 2, 3, 4, 5) - roots = {dsu.find(i) for i in [1, 2, 3, 4, 5]} - assert len(roots) == 1 - - -def test_union_multiple_elements_partial_overlap(dsu): - dsu.union(1, 2) - dsu.union(3, 4) - dsu.union(2, 3, 5) - - roots = {dsu.find(i) for i in [1, 2, 3, 4, 5]} - assert len(roots) == 1 - - -def test_union_multiple_elements_with_existing_connections(dsu): - dsu.union(1, 2) - dsu.union(2, 3) - dsu.union(3, 4, 5, 6) - - roots = {dsu.find(i) for i in [1, 2, 3, 4, 5, 6]} - assert len(roots) == 1 - - -def test_union_single_element(dsu): - assert dsu.union(1) is False - - -def test_union_no_elements(dsu): - assert dsu.union() is False - - -def test_union_self(dsu): - assert dsu.union(3, 3) is False - assert dsu.find(3) == 3 - - -def test_transitivity(dsu): - dsu.union(1, 2) - dsu.union(2, 3) - dsu.union(3, 4) - assert dsu.find(1) == dsu.find(4) - - -def test_disjoint_sets_remain_separate(dsu): - dsu.union(1, 2) - dsu.union(3, 4) - assert dsu.find(1) != dsu.find(3) - - -def test_randomized_unions(): - n = 50 - dsu = DisjointSet(n) - - groups = [{i} for i in range(n)] - - def find_group(x): - for g in groups: - if x in g: - return g - - for _ in range(100): - elems = random.sample(range(n), random.randint(2, 5)) - dsu.union(*elems) - - # merge ground-truth groups - merged = set() - for e in elems: - merged |= find_group(e) - - groups = [g for g in groups if g.isdisjoint(merged)] - groups.append(merged) - - # verify structure matches ground truth - for g in groups: - roots = {dsu.find(x) for x in g} - assert len(roots) == 1 - - -def test_path_compression_effect(): - dsu = DisjointSet(6) - dsu.union(0, 1) - dsu.union(1, 2) - dsu.union(2, 3) - dsu.union(3, 4) - - # Trigger compression - root_before = dsu.find(4) - root_after = dsu.find(4) - - assert root_before == root_after diff --git a/tests/test_ops_monoid.py b/tests/test_ops_monoid.py index fcd72f064..4d243ca14 100644 --- a/tests/test_ops_monoid.py +++ b/tests/test_ops_monoid.py @@ -451,7 +451,10 @@ def test_reduce_independent_2(backend: Backend): lhs = Sum.reduce(Product.plus(a(), b(), f(b(), c())), {a: A(), b: B(), c: C()}) rhs = Product.plus( Sum.reduce(Product.plus(a()), {a: A()}), - Sum.reduce(Product.plus(b(), f(b(), c())), {b: B(), c: C()}), + Sum.reduce( + Product.plus(b(), Sum.reduce(Product.plus(f(b(), c())), {c: C()})), + {b: B()}, + ), ) backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceFactorization()) @@ -489,7 +492,77 @@ def test_reduce_independent_4(backend: Backend): rhs = Product.plus( 7, Sum.reduce(Product.plus(a()), {a: A()}), - Sum.reduce(Product.plus(b(), f(b(), c())), {b: B(), c: C()}), + Sum.reduce( + Product.plus(b(), Sum.reduce(Product.plus(f(b(), c())), {c: C()})), + {b: B()}, + ), + ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceFactorization()) + + +def test_reduce_chain(backend: Backend): + x, y = backend.define_vars("x", "y", ret="scalar") + X, Y = backend.define_vars("X", "Y", ret="stream") + f, h = backend.define_vars("f", "h", arg_types=(backend.scalar_typ,), ret="scalar") + g = backend.define_vars( + "g", arg_types=(backend.scalar_typ, backend.scalar_typ), ret="scalar" + ) + + lhs = Sum.reduce(Product.plus(f(x()), g(x(), y()), h(y())), {x: X(), y: Y()}) + rhs = Sum.reduce( + Product.plus(h(y()), Sum.reduce(Product.plus(f(x()), g(x(), y())), {x: X()})), + {y: Y()}, + ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceFactorization()) + + +@pytest.mark.parametrize("outer,inner", MONOID_PAIRS) +def test_reduce_lift_shared(outer, inner, backend: Backend): + """A stream free in every factor is hoisted into an outer reduce: + Sum.reduce(f(a, c) * g(b, c), {a: A, b: B, c: C}) + = Sum.reduce(Sum.reduce(f(a, c), {a: A}) * Sum.reduce(g(b, c), {b: B}), {c: C}) + """ + a, b, c = backend.define_vars("a", "b", "c", ret="scalar") + A, B, C = backend.define_vars("A", "B", "C", ret="stream") + f, g = backend.define_vars( + "f", "g", arg_types=(backend.scalar_typ, backend.scalar_typ), ret="scalar" + ) + + lhs = outer.reduce(inner.plus(f(a(), c()), g(b(), c())), {a: A(), b: B(), c: C()}) + rhs = outer.reduce( + inner.plus( + outer.reduce(inner.plus(f(a(), c())), {a: A()}), + outer.reduce(inner.plus(g(b(), c())), {b: B()}), + ), + {c: C()}, + ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceFactorization()) + + +@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`` + and ``d = h(c)`` appear in every factor, so both are hoisted.""" + a, b, c, d = backend.define_vars("a", "b", "c", "d", ret="scalar") + A, B, C = backend.define_vars("A", "B", "C", ret="stream") + h = backend.define_vars("h", arg_types=(backend.scalar_typ,), ret="stream") + f, g = backend.define_vars( + "f", + "g", + arg_types=(backend.scalar_typ, backend.scalar_typ, backend.scalar_typ), + ret="scalar", + ) + + lhs = outer.reduce( + inner.plus(f(a(), c(), d()), g(b(), c(), d())), + {a: A(), b: B(), c: C(), d: h(c())}, + ) + rhs = outer.reduce( + inner.plus( + outer.reduce(inner.plus(f(a(), c(), d())), {a: A()}), + outer.reduce(inner.plus(g(b(), c(), d())), {b: B()}), + ), + {c: C(), d: h(c())}, ) backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceFactorization()) From b7cc5434cf1634b46bdbc374e0ee8f7065c72957 Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Thu, 30 Jul 2026 13:03:13 -0400 Subject: [PATCH 08/16] Add weighted einsum implementation (#671) * more precise stream type * add tests for weighted rules * add reduction rule for weighted streams and tests * add test to demo expectation * add numpyro monoid module * add quadrature * add tests * wip * refactor tests * wip * test composition of lifting and weighting * drop numpyro changes * drop unused ops * lint * make weighted a Monoid method * fix typing of jax arrays * change weighted typing to take callable * fix test * fix test * resolve type aliases before dispatching * wip * wip * remove typeof_full * wip * wip * wip * format * refactor test harness * fix behavior of delta terms * add baseline einsum * rework einsum to work on shapes instead of concrete tensors * add einsum benchmark * wip * wip * finish sum/product contraction * allow bind_dims to bind nonexistent named dimensions * wip * add custom partial eval for reductions * working benchmarks * fix infinite loop * eliminate identity indexing when possible * wip * handle getitem where dimensions are created * treat any index with bare ops and slice(None) as canonical * simplify range op and add reduction rules * wip * remove old benchmark code * another try at removing identity gathers * refactor * fix test * lint * clean up comment * fix some test failures * drop sketchy bind_dims rule * drop more type-incompatible plus rules * format * fix reduction issue * drop dimension creating behavior from bind_dims * lint * simplify comment * drop partition * fix docstring * handle negative dimension indexing * fix creation of empty tensors * fully restore previous behavior for missing named dims * reduce any arraylike or named tensor * require at least one jax array to reduce * fix typing test * drop typing test * drop einsum parser in favor of opt_einsum * more agressive factorization that hoists shared streams * reduce nesting * comment * replace with simpler push-based rule * format * drop unused disjoint set * remove unused * push multiple streams instead of one at a time * drop contraction ordering handler * fold BindDimsBindDims into default behavior * handle Sum.reduce instead of Monoid.reduce * wip * wip * hacks * extract contraction heuristic * lint * fix test * use a named dimension einsum for contractions * lint * drop custom arange op * wip * simplify by targetting delta rules * wip * fixes * fixes * lint * drop unused * pick up constants but not rest of module * lint --- effectful/handlers/jax/_handlers.py | 123 +++++- effectful/handlers/jax/_terms.py | 71 +++- effectful/handlers/jax/monoid.py | 464 +++++++++++++++-------- effectful/handlers/jax/numpy/__init__.py | 27 +- effectful/handlers/jax/scipy/special.py | 4 +- effectful/ops/monoid.py | 202 +++++----- effectful/ops/syntax.py | 8 + pyproject.toml | 5 +- tests/test_handlers_jax_monoid.py | 431 ++++++++++++++++----- tests/test_ops_monoid.py | 38 +- 10 files changed, 978 insertions(+), 395 deletions(-) diff --git a/effectful/handlers/jax/_handlers.py b/effectful/handlers/jax/_handlers.py index 920aa876e..4c021e9b6 100644 --- a/effectful/handlers/jax/_handlers.py +++ b/effectful/handlers/jax/_handlers.py @@ -1,17 +1,20 @@ import functools import itertools import typing -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Callable, Iterable, Mapping, Sequence from types import EllipsisType from typing import Annotated +from opt_einsum import get_symbol +from opt_einsum.parser import parse_einsum_input + try: import jax import jax.numpy as jnp except ImportError: raise ImportError("JAX is required to use effectful.handlers.jax") -from effectful.ops.semantics import apply, evaluate, fvsof, typeof +from effectful.ops.semantics import apply, evaluate, fvsof, fwd, typeof from effectful.ops.syntax import ( ConstructorOperation, PureInterpretation, @@ -83,11 +86,17 @@ def _getitem(arr, index): if not is_eager_array(term_arr): return functools.reduce(_merge, arg_sizes, {}) - sizes = ( - {k.op: term_arr.shape[i]} - for i, k in enumerate(term_index) - if isinstance(k, Term) and len(k.args) == 0 and len(k.kwargs) == 0 - ) + sizes = [] + i = 0 + for k in term_index: + if isinstance(k, Term) and len(k.args) == 0 and len(k.kwargs) == 0: + sizes.append({k.op: term_arr.shape[i]}) + + # None in the index expression creates a new dimension, which does + # not appear in the shape of the indexed tensor + if k is not None: + i += 1 + return functools.reduce(_merge, itertools.chain(arg_sizes, sizes), {}) return ( @@ -140,8 +149,9 @@ def _partial_eval(t: Expr[jax.Array]) -> Expr[jax.Array]: # if any dimension is zero sized, the result is empty if any(size == 0 for size in sized_fvs.values()): - key = tuple(sized_fvs.keys()) - shape = tuple(sized_fvs[k] for k in key) + ops = tuple(sized_fvs.keys()) + key = tuple(k() for k in ops) + shape = tuple(sized_fvs[k] for k in ops) return jax_getitem(jnp.empty(shape), key) def _is_eager(t): @@ -193,7 +203,11 @@ def _jax_op(*args, **kwargs) -> jax.Array: and not isinstance(args[0], Term) and sized_fvs and args[1] - and all(isinstance(k, Term) and k.op in sized_fvs for k in args[1]) + and all( + (isinstance(k, Term) and k.op in sized_fvs) + or (isinstance(k, slice) and k == slice(None)) + for k in args[1] + ) ): raise NotHandled elif sized_fvs and set(sized_fvs.keys()) == fvsof(tm) - {jax_getitem, _jax_op}: @@ -233,6 +247,93 @@ def _jax_op(*args, **kwargs) -> jax.Array: return _jax_op +def _named_dims(term: Expr[jax.Array]) -> tuple[Operation, ...]: + if not (isinstance(term, Term) and term.op == jax_getitem): + return () + index = term.args[1] + assert isinstance(index, Iterable) + return tuple(i.op for i in index if isinstance(i, Term) and not i.args) + + +def _reduce_named(array, axis=None, **kwargs) -> jax.Array: + if axis is None: + return fwd() + + named_dims = _named_dims(array) + if not named_dims: + return fwd() + + bound_arr = bind_dims(array, *named_dims) + + if isinstance(axis, int): + axis = (axis,) + shifted_axis = tuple(a + len(named_dims) if a >= 0 else a for a in axis) + + reduced = fwd(bound_arr, axis=shifted_axis, **kwargs) + return unbind_dims(reduced, *named_dims) + + +def _einsum_named(subscripts, *operands, **kwargs) -> jax.Array: + # only the string-subscripts form is handled; forward the interleaved form + if not isinstance(subscripts, str): + if any(isinstance(x, Term) for x in (subscripts, *operands)): + raise ValueError("Interleaved einsum is not implemented with named tensors") + return jax.numpy.einsum(subscripts, *operands, **kwargs) + + # forward if any operand has a symbolic (Term) shape + if any(isinstance(arr.shape, Term) for arr in operands): + raise NotHandled + + named = [_named_dims(op) for op in operands] + + # normalize: expand ellipses and make the output explicit, using the + # positional shapes (shapes=True avoids materializing the operands) + shapes = [op.shape for op in operands] + in_part, out_part, _ = parse_einsum_input([subscripts, *shapes], shapes=True) + in_specs = in_part.split(",") + assert len(in_specs) == len(operands) + + # fresh symbols for named dims, avoiding every symbol already in use; + # get_symbol gives an effectively unlimited supply (spills into unicode) + used = {c for c in (in_part + out_part) if c not in ",->"} + counter = itertools.count() + + def next_symbol(): + while True: + s = get_symbol(next(counter)) + if s not in used: + used.add(s) + return s + + # assign a letter per unique named dim; shared names reuse the same letter + # so einsum aligns them as batch dims rather than contracting + letter_of, order = {}, [] + for dims in named: + for d in dims: + if d not in letter_of: + letter_of[d] = next_symbol() + order.append(d) + + # bind named dims to leading positional axes and prepend their letters + bound, new_in_specs = [], [] + for op, dims, spec in zip(operands, named, in_specs): + bound.append(bind_dims(op, *dims) if dims else op) + new_in_specs.append("".join(letter_of[d] for d in dims) + spec) + + # add every named dim to the front of the output as passthrough + out_prefix = "".join(letter_of[d] for d in order) + new_subscripts = ",".join(new_in_specs) + "->" + out_prefix + out_part + + result = jax.numpy.einsum(new_subscripts, *bound, **kwargs) + + # unbind: leading axes correspond to `order`, reindex them back to named + reindexed = jax_getitem( + result, + tuple(d() for d in order) + tuple(slice(None) for _ in range(len(out_part))), + ) + return reindexed + + @_register_jax_op def jax_getitem(x: jax.Array, key: tuple[IndexElement, ...]) -> jax.Array: """Operation for indexing an array. Unlike the standard __getitem__ method, @@ -266,6 +367,8 @@ def bind_dims[T, A, B]( >>> bind_dims(t, b, a).shape (3, 2) """ + if isinstance(value, Term) and value.op == bind_dims: + return bind_dims(value.args[0], *(names + tuple(value.args[1:]))) if jax.tree_util.treedef_is_leaf(jax.tree.structure(value)): return __dispatch(typeof(value))(value, *names) return jax.tree.map(lambda v: bind_dims(v, *names), value) diff --git a/effectful/handlers/jax/_terms.py b/effectful/handlers/jax/_terms.py index 05a5390e7..53c4c094f 100644 --- a/effectful/handlers/jax/_terms.py +++ b/effectful/handlers/jax/_terms.py @@ -467,31 +467,66 @@ def _bind_dims_array(t: jax.Array, *args: Operation[[], jax.Array]) -> jax.Array array = t.args[0] dims = t.args[1] assert isinstance(dims, Sequence) + ndim = len(array.shape) # ensure that the order is a subset of the named dimensions order_set = set(args) if not order_set <= set(a.op for a in dims if isinstance(a, Term)): raise NotHandled - # permute the inner array so that the leading dimensions are in the order - # specified and the trailing dimensions are the remaining named dimensions - # (or slices) - reindex_dims = [ - i - for i, o in enumerate(dims) - if not isinstance(o, Term) or o.op not in order_set - ] - dim_ops = [a.op if isinstance(a, Term) else None for a in dims] - perm = ( - [dim_ops.index(o) for o in args] - + reindex_dims - + list(range(len(dims), len(array.shape))) + def axis_op(ax: int) -> Operation | None: + """The named op of a bare index term at axis ``ax``, else ``None``.""" + if ax < len(dims): + d = dims[ax] + if isinstance(d, Term) and not d.args and not d.kwargs: + return d.op + return None + + # Assign an einsum id to every axis of ``array``. Axes that share a named op + # get the *same* id — a repeated op (e.g. ``arr[i(), i()]``) ties its axes + # together, which einsum reads as a diagonal. Every other axis (slices, ints, + # fancy indices, compound terms, and trailing positional axes) gets a unique + # id, so einsum simply carries it through to be reindexed below. + op_ids: dict[Operation, int] = {} + in_ids: list[int] = [] + next_id = 0 + for ax in range(ndim): + op = axis_op(ax) + if op is not None: + if op not in op_ids: + op_ids[op] = next_id + next_id += 1 + in_ids.append(op_ids[op]) + else: + in_ids.append(next_id) + next_id += 1 + + # Output order: bound args that actually appear (in the requested order, + # deduplicated by the diagonal merge), then every remaining axis in + # first-appearance order. einsum does the permutation and the diagonals; with + # all distinct ids retained in the output it performs no reduction. + present_arg_ids = [op_ids[o] for o in args if o in op_ids] + seen = set(present_arg_ids) + rest_ids: list[int] = [] + for i in in_ids: + if i not in seen: + seen.add(i) + rest_ids.append(i) + + array = jnp.einsum(array, in_ids, present_arg_ids + rest_ids) + + # Re-apply the original index for each carried axis and re-name unbound op + # axes. Trailing positional axes (first appearance beyond ``dims``) are left + # for jax_getitem to carry implicitly. + first_pos: dict[int, int] = {} + for ax, i in enumerate(in_ids): + first_pos.setdefault(i, ax) + + index_expr = (slice(None),) * len(present_arg_ids) + tuple( + dims[first_pos[i]] if first_pos[i] < len(dims) else slice(None) + for i in rest_ids ) - array = jnp.transpose(array, perm) - reindexed = jax_getitem( - array, (slice(None),) * len(args) + tuple(dims[i] for i in reindex_dims) - ) - return reindexed + return jax_getitem(array, index_expr) @unbind_dims.register(jax.Array) # type: ignore diff --git a/effectful/handlers/jax/monoid.py b/effectful/handlers/jax/monoid.py index 3f6273be3..fbd845d8b 100644 --- a/effectful/handlers/jax/monoid.py +++ b/effectful/handlers/jax/monoid.py @@ -1,11 +1,16 @@ import functools +import logging import typing -from collections.abc import Iterable +from typing import Protocol import jax +import jax.core +import opt_einsum +from opt_einsum import get_symbol import effectful.handlers.jax.numpy as jnp -from effectful.handlers.jax import bind_dims, unbind_dims +from effectful.handlers.jax import bind_dims, jax_getitem, unbind_dims +from effectful.handlers.jax._handlers import is_eager_array from effectful.handlers.jax.scipy.special import logsumexp from effectful.ops.monoid import ( CartesianProduct, @@ -16,13 +21,16 @@ Product, Streams, Sum, + _is_monoid_plus, + choose_contraction, distributes_over, - outer_stream, ) from effectful.ops.semantics import evaluate, fvsof, fwd, handler, typeof from effectful.ops.syntax import ObjectInterpretation, deffn, implements from effectful.ops.types import Interpretation, NotHandled, Operation, Term +logger = logging.getLogger(__name__) + def cartesian_prod(x, y): if x.ndim == 1: @@ -46,16 +54,39 @@ def cartesian_prod(x, y): def _jax_args(args): """True iff ``args`` is non-empty and every arg is a concrete - :class:`jax.Array` (no Terms). + :class:`jax.typing.ArrayLike` or named tensor. At least one argument must be + a jax-related type. + """ - typs = (typeof(a) for a in args) return ( bool(args) - and any(issubclass(t, jax.Array) for t in typs) - and all(issubclass(t, jax.typing.ArrayLike) for t in typs) + and all(is_eager_array(a) or isinstance(a, jax.typing.ArrayLike) for a in args) + and any(is_eager_array(a) or isinstance(a, jax.Array) for a in args) ) +class PlusJaxUpcast(ObjectInterpretation): + @implements(Monoid.plus) + def plus(self, monoid, *args): + arg_types = [typeof(a) for a in args] + + def _is_jax(t): + return issubclass(t, jax.Array | jax.core.Tracer) + + # exists array valued and non-array-valued args + if any(_is_jax(t) for t in arg_types) and any( + not _is_jax(t) for t in arg_types + ): + return monoid.plus( + *( + a if _is_jax(t) else jnp.asarray(a) + for (a, t) in zip(args, arg_types, strict=True) + ) + ) + + return fwd() + + class SumPlusJax(ObjectInterpretation): @implements(Sum.plus) def plus(self, *args): @@ -124,125 +155,144 @@ def plus(self, *args): return result -ARRAY_REDUCTORS = { - Sum: jnp.sum, - Product: jnp.prod, - Min: jnp.min, - Max: jnp.max, - LogSumExp: logsumexp, -} +class ReduceArrayGather(ObjectInterpretation): + """M.reduce(body, {k: a} ∪ S) ≡ M.reduce(body[k := a[k']], {k': range(a.shape[0])} ∪ S)""" - -class ArrayReduce(ObjectInterpretation): @implements(Monoid.reduce) def reduce(self, monoid, body, streams): - if monoid not in ARRAY_REDUCTORS or typeof(body) is not jax.Array: + if typeof(body) is not jax.Array: return fwd() - if not streams: - return monoid.identity - reductor = ARRAY_REDUCTORS[monoid] - index = Operation.define(jax.Array) - for stream_key, stream_body, streams_tail in outer_stream(streams): - if not issubclass(typeof(stream_body), jax.Array): - continue + if isinstance(body, Term) and body.op is delta: + return fwd() - if stream_key in fvsof(body): - with handler({stream_key: deffn(unbind_dims(stream_body, index))}): - eval_body = evaluate(body) - eval_streams_tail = evaluate(streams_tail) - assert isinstance(eval_streams_tail, dict) - reduce_tail = ( - monoid.reduce(eval_body, eval_streams_tail) - if len(eval_streams_tail) > 0 - else eval_body - ) - return reductor(bind_dims(reduce_tail, index), axis=0) + body_fvs = fvsof(body) + stream_keys = set(streams) + + body_subst = {} + streams_subst = {} + range_streams = {} + progress = False + for k, v in streams.items(): + if is_eager_array(v) and k in body_fvs and not (fvsof(v) & stream_keys): + kk = Operation.define(k) + body_subst[k] = deffn(unbind_dims(v, kk)) + streams_subst[k] = kk + range_streams[kk] = range(v.shape[0]) + progress = True else: - # TODO: In this case, the stream is unused in the body. The body - # should be multiplied by the length of the stream. The current - # behavior is not efficient. - return fwd() + range_streams[k] = v - return fwd() + if not progress: + return fwd() + subst_body = handler(body_subst)(evaluate)(body) + subst_streams = handler(streams_subst)(evaluate)(range_streams) + return monoid.reduce(subst_body, subst_streams) -@Operation.define -def delta(_index: tuple[int, ...], _weight: jax.Array) -> jax.Array: - raise NotHandled +class Reductor(Protocol): + def __call__( + self, arr: jax.Array, axis: int | tuple[int, ...] | None = None + ) -> jax.Array: ... -py_range = range +ARRAY_REDUCTORS: dict[Monoid, Reductor] = {} +for monoid, func in [ + (Sum, jnp.sum), + (Product, jnp.prod), + (Min, jnp.min), + (Max, jnp.max), +]: + assert isinstance(monoid, Monoid) + assert callable(func) + ARRAY_REDUCTORS[monoid] = functools.partial(func, initial=monoid.identity) -@Operation.define -def range(*args: int) -> Iterable[jax.Array]: - raise NotHandled +ARRAY_REDUCTORS[LogSumExp] = logsumexp -def _range_start(term: Term): - assert term.op == range - if len(term.args) < 2: - return 0 - return term.args[0] +class ReduceArray(ObjectInterpretation): + """Reduce an array body over range streams.""" + + @implements(Monoid.reduce) + def reduce(self, monoid, body, streams): + reductor = ARRAY_REDUCTORS.get(monoid, None) + if reductor is None: + return fwd() + + if typeof(body) is not jax.Array: + return fwd() + + pos_dims = {} + if isinstance(body, Term): + if body.op == delta: + pos_dims = { + d.op + for d in body.args[0] + if isinstance(d, Term) and d.op in streams + } + elif _is_monoid_plus(body.op) and distributes_over( + body.op.__self__, monoid + ): + # delegate to factorization + return fwd() + + body_fvs = fvsof(body) + used = { + k + for k, v in streams.items() + if k in body_fvs and k not in pos_dims and isinstance(v, range) + } + if not used: + return fwd() + + delta_key = tuple(k() for k in streams if k in used) + arr = monoid.reduce(delta(delta_key, body), streams) + reduced_body = reductor(arr, axis=tuple(range(len(used)))) + return reduced_body + + +@Operation.define +def delta(_index: tuple[int, ...], _weight: jax.Array) -> jax.Array: + raise NotHandled def _range_stop(term: Term): - assert term.op == range + assert term.op == jnp.arange + if "stop" in term.kwargs: + return term.kwargs["stop"] if len(term.args) < 2: return term.args[0] return term.args[1] -def _range_step(term: Term): - assert term.op == range - if len(term.args) < 3: - return 1 - return term.args[2] +class DeltaEmpty(ObjectInterpretation): + """delta((), weight) ≡ weight""" + + @implements(delta) + def _(self, index, weight): + if not index: + return weight + return fwd() -def _is_simple_range(term: Term) -> bool: - if term.op != range: - return False +class DeltaFusion(ObjectInterpretation): + """delta(i1, delta(i2, weight)) ≡ delta(i1 ++ i2, weight)""" - start = _range_start(term) - step = _range_step(term) - return ( - not isinstance(start, Term) - and start == 0 - and not isinstance(step, Term) - and step == 1 - ) + @implements(delta) + def _(self, index, weight): + if isinstance(weight, Term) and weight.op == delta: + return delta(index + weight.args[0], weight.args[1]) + return fwd() -class ReduceDeltaIndependent(ObjectInterpretation): +class ReduceDeltaSimpleRange(ObjectInterpretation): """Eliminate a Delta that has independent, dense index arguments. - reduce(M, streams, delta((), body)) ≡ reduce(M, streams, body) - reduce(M, streams ∪ {v: range(N)}, delta(idx' ++ (v(),), body)) + reduce(M, streams ∪ {v: range(N)}, delta((v(),) ++ idx', body)) ═══════════════════════════════════════════════════════════════════════════ - reduce(M, streams, delta(idx', bind_dims(body[v() := unbind_dims(streams[v], fv)], fv))) - - Not yet supported: - - - **Strided index streams** (``range(0, N, k)`` for ``k != 1``): the - premise ``_is_simple_range`` requires ``start == 0`` and ``step == 1``. - A strided extension would substitute ``v() := unbind_dims(jnp.arange( - start, stop, step), fv)`` and otherwise follow the same shape — the - change is purely in the recognised range form, the bind/unbind cycle - below is unchanged. - - **Non-zero start** (``range(a, b, 1)`` with ``a != 0``): same template - as the strided case; only the recognised range form changes. - - **Non-bare index expressions** (``delta((2*v(),), w)``, - ``delta((f(v()),), w)``, etc.): currently requires the final index - entry to be a bare call ``v()`` of a stream var op. Generalizing to - arbitrary index expressions is a scatter, not a bind: materialize the - index expression and the weight separately over ``v``, then - ``jnp.zeros(N).at[indices].set(values)`` (for Sum; analogous for - other monoids using ``.add``/``.min``/``.max``/...). This is a - different leaf operation from ``bind_dims`` and warrants a sibling - rule rather than an extension of this one. + bind_dims(reduce(M, streams, delta(idx', body[v() := unbind_dims(streams[v], fv)])), fv) """ @implements(Monoid.reduce) @@ -250,31 +300,76 @@ def _(self, monoid: Monoid, body, streams: Streams): if not (isinstance(body, Term) and body.op == delta): return fwd() - indices, weight = body.args - assert isinstance(indices, tuple) + index, weight = body.args + assert isinstance(index, tuple) - if not indices: - return monoid.reduce(weight, streams) + if not index: + return fwd() - head_indices, tail_index = indices[:-1], indices[-1] - if not (isinstance(tail_index, Term) and tail_index.op in streams): + head_index, tail_index = index[0], index[1:] + if not (isinstance(head_index, Term) and head_index.op in streams): return fwd() - tail_op: Operation = tail_index.op - tail_stream = streams[tail_op] - if not (isinstance(tail_stream, Term) and _is_simple_range(tail_stream)): + head_op: Operation = head_index.op + head_stream = streams[head_op] + if not ( + isinstance(head_stream, range) + and head_stream.start == 0 + and head_stream.step == 1 + ): return fwd() - fresh_op = Operation.define(tail_op) - indices = jnp.arange(_range_stop(tail_stream)) - if isinstance(indices, jax.Array) and len(indices) == 0: - return monoid.identity + tail_streams = {k: v for (k, v) in streams.items() if k != head_op} + + # peel the head index: substitute it into the weight (slicing direct + # uses, materializing the rest) along a fresh named dim, but bind that + # dim only *after* the surrounding reduce -- see the class docstring. + + fresh_op = Operation.define(head_op) + + def _jax_getitem(arr, index): + inner_index, outer_index = [], [] + progress = False + for i in index: + if isinstance(i, Term) and i.op == head_op: + inner_index.append( + slice(head_stream.start, head_stream.stop, head_stream.step) + ) + outer_index.append(fresh_op()) + progress = True + else: + inner_index.append(slice(None)) + outer_index.append(i) + if progress: + return jax_getitem(jax_getitem(arr, inner_index), outer_index) + return fwd(arr, index) + + slice_subst = typing.cast(Interpretation, {jax_getitem: _jax_getitem}) + sliced_weight = handler(slice_subst)(evaluate)(weight) + sliced_streams = handler(slice_subst)(evaluate)(tail_streams) + + gather_subst = typing.cast( + Interpretation, + { + head_op: deffn( + unbind_dims( + jnp.arange( + head_stream.start, head_stream.stop, head_stream.step + ), + fresh_op, + ) + ) + }, + ) + gathered_weight = handler(gather_subst)(evaluate)(sliced_weight) + gathered_streams = handler(gather_subst)(evaluate)(sliced_streams) - fresh_stream = unbind_dims(indices, fresh_op) - subst_intp = typing.cast(Interpretation, {tail_op: deffn(fresh_stream)}) - fresh_body = bind_dims(handler(subst_intp)(evaluate)(weight), fresh_op) - fresh_streams = {k: v for (k, v) in streams.items() if k != tail_op} - return monoid.reduce(delta(head_indices, fresh_body), fresh_streams) + inner = ( + monoid.reduce(delta(tail_index, gathered_weight), gathered_streams) + if gathered_streams + else gathered_weight + ) + return bind_dims(inner, fresh_op) class ReduceDependentRangeMask(ObjectInterpretation): @@ -324,15 +419,16 @@ def _(self, monoid: Monoid, body, streams: Streams): simple_ranges = { k: v for (k, v) in streams.items() - if isinstance(v, Term) and _is_simple_range(v) + if isinstance(v, range) and v.start == 0 and v.step == 1 } for u, u_stream in simple_ranges.items(): if fvsof(u_stream) & stream_vars: continue - for v, v_stream in simple_ranges.items(): + for v, v_stream in streams.items(): if ( isinstance(v_stream, Term) + and v_stream.op == jnp.arange and isinstance(_range_stop(v_stream), Term) and _range_stop(v_stream).op == u ): @@ -355,38 +451,6 @@ def _(self, monoid: Monoid, body, streams: Streams): return fwd() -class ReduceRange(ObjectInterpretation): - """Replace concrete-range stream values with materialized ``jnp.arange``. - - reduce(M, streams ∪ {v: range(a, b, s)}, body) - ≡ reduce(M, streams ∪ {v: jnp.arange(a, b, s)}, body) - - when ``a``, ``b``, ``s`` are concrete and ``body`` is not a delta term. - Delegates the actual reduction to whichever handler picks up the - materialized ``jax.Array`` streams. - """ - - @implements(Monoid.reduce) - def _(self, monoid: Monoid, body, streams: Streams): - if isinstance(body, Term) and body.op == delta: - return fwd() - - new_streams: dict = {} - any_replaced = False - for k, v in streams.items(): - if isinstance(v, Term) and v.op == range: - new_streams[k] = jnp.arange( - _range_start(v), _range_stop(v), _range_step(v) - ) - any_replaced = True - else: - new_streams[k] = v - - if not any_replaced: - return fwd() - return monoid.reduce(body, new_streams) - - # Cross-cutting delta rules not yet implemented: # # - **Delta-commuting** (DC-hoist): for any pure op ``f`` (no Scoped binders @@ -406,23 +470,115 @@ def _(self, monoid: Monoid, body, streams: Streams): # a subsequence of ``idx_b`` (or vice versa). Refuse to fire when neither # is a subsequence of the other, since that would silently insert an # outer-product broadcast. -# -# - **Empty-domain detection at the term level**: currently size-0 named -# dims must be resolved by leaf consumers (``bind_dims``, reductors with -# ``initial=monoid.identity``). The empty-domain check is intentionally -# NOT a rule on its own — rewrites stay size-polymorphic and leaf ops -# carry the burden. See the conversation in monoid.py's history for why. + + +class ContractLongestArrayStream(ObjectInterpretation): + @implements(choose_contraction) + def _(self, factors, streams): + lengths = { + k: v.shape[0] if isinstance(v, jax.Array) and v.shape else 0 + for (k, v) in streams.items() + } + longest = max(lengths.values()) + return fwd( + factors, {k: v for (k, v) in streams.items() if lengths[k] == longest} + ) + + +class ReduceSumProductContraction(ObjectInterpretation): + """Fast-path a sum-of-products contraction.""" + + @implements(Sum.reduce) + def _(self, body, streams: Streams): + if not ( + isinstance(body, Term) + and _is_monoid_plus(body.op) + and body.op.__self__ is Product + ): + return fwd() + + factors = body.args + if len(factors) != 2 or not all( + issubclass(typeof(f), jax.Array) for f in factors + ): + return fwd() + + (lhs, rhs) = factors + stream_vars = set(streams.keys()) + + # a fully factored reduce only has streams that are used by all factors + shared = fvsof(lhs) & fvsof(rhs) & stream_vars + if shared != stream_vars: + return fwd() + + if not all(isinstance(v, range) for v in streams.values()): + return fwd() + + # create leading reduction dimensions + delta_key = tuple(k() for k in streams) + pos_lhs = Sum.reduce(delta(delta_key, lhs), streams) + pos_rhs = Sum.reduce(delta(delta_key, rhs), streams) + + dims = "".join(get_symbol(i) for i in range(len(streams))) + contraction = jnp.einsum(f"{dims}...,{dims}...->...", pos_lhs, pos_rhs) + return contraction + + +@jax.jit(static_argnums=(0,)) +def einsum(subscripts: str, /, *operands: jax.Array) -> jax.Array: + """Evaluate an einsum expression using monoid reductions.""" + if not operands: + raise ValueError("einsum requires at least one operand") + + in_spec, out_spec, _ = opt_einsum.parser.parse_einsum_input( + [subscripts, *(op.shape for op in operands)], shapes=True + ) + in_specs = in_spec.split(",") + + all_letters = set(out_spec) | {c for s in in_specs for c in s} + ops = {c: Operation.define(jax.Array, name=c) for c in all_letters} + + sizes: dict[str, int] = {} + for spec, op in zip(in_specs, operands, strict=True): + for l, s in zip(spec, op.shape, strict=True): + if l in sizes and sizes[l] != s: + raise ValueError(f"Dimension {l} given sizes {s} and {sizes[l]}") + else: + sizes[l] = s + for c in out_spec: + if c not in sizes: + raise ValueError(f"einsum: output index {c!r} not present in any input") + + arrays = [Operation.define(jax.Array) for _ in operands] + factors = [ + unbind_dims(arr(), *(ops[c] for c in spec)) + for arr, spec in zip(arrays, in_specs, strict=True) + ] + body = Product.plus(*factors) + + out_tuple = tuple(ops[c]() for c in out_spec) + streams = {op: range(sizes[c]) for c, op in ops.items()} + with handler(NormalizeIntp): + norm = deffn(Sum.reduce(delta(out_tuple, body), streams), *arrays) + result = norm(*operands) + assert isinstance(result, jax.Array) + return result NormalizeIntp.extend( - ArrayReduce(), - ReduceRange(), - ReduceDeltaIndependent(), + ReduceArray(), + ReduceSumProductContraction(), + ReduceArrayGather(), + ReduceDeltaSimpleRange(), ReduceDependentRangeMask(), + DeltaEmpty(), + DeltaFusion(), SumPlusJax(), ProductPlusJax(), MinPlusJax(), MaxPlusJax(), LogSumExpPlusJax(), CartesianProductPlusJax(), + ContractLongestArrayStream(), + PlusJaxUpcast(), ) diff --git a/effectful/handlers/jax/numpy/__init__.py b/effectful/handlers/jax/numpy/__init__.py index cc20d7498..9556bdacf 100644 --- a/effectful/handlers/jax/numpy/__init__.py +++ b/effectful/handlers/jax/numpy/__init__.py @@ -3,9 +3,17 @@ import jax.numpy -from .._handlers import _register_jax_op, _register_jax_op_no_partial_eval - -_no_overload = ["array", "asarray"] +from effectful.handlers.jax._handlers import ( + _einsum_named, + _reduce_named, + _register_jax_op, + _register_jax_op_no_partial_eval, +) +from effectful.ops.semantics import handler +from effectful.ops.types import Operation + +_NO_OVERLOAD = ["array", "asarray"] +_REDUCTION = ["sum", "prod", "min", "max", "any", "all", "mean", "argmax"] for name, op in jax.numpy.__dict__.items(): if isinstance(op, types.ModuleType): @@ -19,18 +27,25 @@ if name == "__getattr__": continue - elif name in _no_overload: + elif name in _NO_OVERLOAD: globals()[name] = _register_jax_op_no_partial_eval(op) else: globals()[name] = _register_jax_op(op) jax_op = ( _register_jax_op_no_partial_eval(op) - if name in _no_overload + if name in _NO_OVERLOAD else _register_jax_op(op) ) globals()[name] = jax_op +for name in _REDUCTION: + op = globals()[name] + globals()[name] = handler({op: _reduce_named})(op) + + +einsum = Operation.define(_einsum_named) + # Tell mypy about our wrapped functions. if TYPE_CHECKING: - from jax.numpy import * # noqa: F403 + from jax.numpy import * # type: ignore[assignment] # noqa: F403 diff --git a/effectful/handlers/jax/scipy/special.py b/effectful/handlers/jax/scipy/special.py index afe1334b8..67b99621a 100644 --- a/effectful/handlers/jax/scipy/special.py +++ b/effectful/handlers/jax/scipy/special.py @@ -2,9 +2,11 @@ import jax.scipy.special -from effectful.handlers.jax._handlers import _register_jax_op +from effectful.handlers.jax._handlers import _reduce_named, _register_jax_op +from effectful.ops.semantics import handler logsumexp = _register_jax_op(jax.scipy.special.logsumexp) +logsumexp = handler({logsumexp: _reduce_named})(logsumexp) # Tell mypy about our wrapped functions. if TYPE_CHECKING: diff --git a/effectful/ops/monoid.py b/effectful/ops/monoid.py index 5f342f25b..42180e912 100644 --- a/effectful/ops/monoid.py +++ b/effectful/ops/monoid.py @@ -4,34 +4,22 @@ import operator import typing from collections import Counter, UserDict, defaultdict -from collections.abc import Callable, Generator, Iterable, Mapping +from collections.abc import Callable, Generator, Iterable, Mapping, Sequence from dataclasses import dataclass from graphlib import TopologicalSorter from typing import Annotated, Any -from effectful.ops.semantics import ( - coproduct, - evaluate, - fvsof, - fwd, - handler, - typeof, -) +from effectful.ops.semantics import coproduct, evaluate, fvsof, fwd, handler, typeof from effectful.ops.syntax import ( ObjectInterpretation, Scoped, + defdata, deffn, implements, syntactic_eq, syntactic_hash, ) -from effectful.ops.types import ( - Expr, - Interpretation, - NotHandled, - Operation, - Term, -) +from effectful.ops.types import Expr, Interpretation, NotHandled, Operation, Term type Stream[T] = Iterable[T] @@ -125,16 +113,21 @@ def __eq__(self, other): def __hash__(self): return hash(id(self)) - # the weak typing allows us to write monoid.plus(monoid.identity, ) - # and monoid.plus(monoid.identity, ) @Operation.define - def plus(self, *args: Any) -> Any: + def plus(self, *args: W) -> W: """Monoid addition. Handlers supply per-monoid and broadcasting - behavior; the default rule only handles empty / Term cases. + behavior; the default rule only handles identity and zero cases (for + monoids that have a zero). + """ - if not args: - return self.identity - raise NotHandled + if hasattr(self, "zero") and any(a is self.zero for a in args): + return self.zero + + nonident_args = [a for a in args if a is not self.identity] + if len(nonident_args) != len(args): + return self.plus(*nonident_args) + + return defdata(self.plus, *nonident_args) # type: ignore[return-value] @Operation.define def reduce[A, B, U: Body]( @@ -146,24 +139,6 @@ def reduce[A, B, U: Body]( broadcasting behavior; the default rule only handles the empty-stream case. """ - for stream_key, stream_body, streams_tail in outer_stream(streams): - if isinstance(stream_body, Term): - continue - stream_values_iter = iter(stream_body) - - # if we iterate and get a term instead of a real iterator, skip - if isinstance(stream_values_iter, Term): - continue - - new_reduces = [] - for stream_val in stream_values_iter: - with handler({stream_key: deffn(stream_val)}): - eval_args = evaluate((body, streams_tail)) - assert isinstance(eval_args, tuple) - new_reduces.append( - self.reduce(*eval_args) if streams_tail else eval_args[0] - ) - return self.plus(*new_reduces) raise NotHandled @Operation.define @@ -268,16 +243,6 @@ def plus(self, _, *args): return fwd() -class PlusIdentity(ObjectInterpretation): - """x₁ + ... + 0 + ... + xₙ = x₁ + ... + xₙ""" - - @implements(Monoid.plus) - def plus(self, monoid, *args): - if any(x is monoid.identity for x in args): - return monoid.plus(*(x for x in args if x is not monoid.identity)) - return fwd() - - class PlusAssoc(ObjectInterpretation): """x + (y + z) = (x + y) + z = x + y + z""" @@ -338,18 +303,6 @@ def plus(self, monoid: Monoid, *args): return fwd() -class PlusZero(ObjectInterpretation): - """x₁ * ... * 0 * ... * xₙ = 0""" - - @implements(Monoid.plus) - def plus(self, monoid, *args): - if not (isinstance(monoid, MonoidWithZero)): - return fwd() - if any(x is monoid.zero for x in args): - return monoid.zero - return fwd() - - class PlusConsecutiveDups(ObjectInterpretation): """x ⊕ x ⊕ y = x ⊕ y""" @@ -397,15 +350,30 @@ def plus(self, monoid, *args): return fwd() -class ReduceNoStreams(ObjectInterpretation): - """Implements the identity - reduce(R, ∅, body) = 0 - """ - +class ReducePartial(ObjectInterpretation): @implements(Monoid.reduce) - def reduce(self, monoid, _, streams): - if len(streams) == 0: + def _(self, monoid, body, streams): + if not streams: return monoid.identity + + for stream_key, stream_body, streams_tail in outer_stream(streams): + if isinstance(stream_body, Term): + continue + stream_values_iter = iter(stream_body) + + # if we iterate and get a term instead of a real iterator, skip + if isinstance(stream_values_iter, Term): + continue + + new_reduces = [] + for stream_val in stream_values_iter: + with handler({stream_key: deffn(stream_val)}): + eval_args = evaluate((body, streams_tail)) + assert isinstance(eval_args, tuple) + new_reduces.append( + monoid.reduce(*eval_args) if streams_tail else eval_args[0] + ) + return monoid.plus(*new_reduces) return fwd() @@ -435,6 +403,30 @@ def reduce(self, monoid, body, streams): return fwd() +@Operation.define +def choose_contraction(factors: Sequence[Any], streams: Streams) -> Operation: + """Used by `ReduceFactorization` to choose a contraction when there is + ambiguity. Takes the factors and streams that are eligible for contraction + (innermost and non-universal). + + The default behavior is to return the first support-minimal stream in the + streams dictionary. + + """ + assert len(streams) > 0 + + factors = [(a, fvsof(a)) for a in factors] + support: dict = { + k: frozenset(i for i, (_, fvs) in enumerate(factors) if k in fvs) + for k in streams + } + for v, f_v in support.items(): + if any(u_sup < f_v for u, u_sup in support.items() if u is not v): + continue + return v + assert False, "expected at least one subset-minimal stream" + + class ReduceFactorization(ObjectInterpretation): """reduce(⊗(F_v ∪ F_rest), {v} ∪ S) = reduce(⊗F_rest ⊗ reduce(⊗F_v, {v}), S) @@ -459,39 +451,40 @@ def reduce(self, monoid, body, streams): # candidates: innermost-eligible (no remaining stream depends on v), # non-universal (some factor doesn't mention v) - support: dict = {} - for v in streams: - if any(v in fvsof(s) for k, s in streams.items() if k is not v): + eligible = {} + for k, v in streams.items(): + if any(k in fvsof(vv) for kk, vv in streams.items() if k is not kk): continue - f_v = frozenset(i for i, (_, fvs) in enumerate(factors) if v in fvs) - if len(f_v) == len(factors): + if len({i for i, (_, fvs) in enumerate(factors) if k in fvs}) == len( + factors + ): continue # v is universal: leave it in the outer core - support[v] = f_v - - # eliminate a variable with subset-minimal factor support - # (leaves-first; canonical on hierarchical/laminar supports) - inner_stream = None - inner_factor_ids = None - for v, f_v in support.items(): - if any(u_sup < f_v for u, u_sup in support.items() if u is not v): - continue - inner_stream = v - inner_factor_ids = f_v - break + eligible[k] = v - if not inner_stream or not inner_factor_ids: + if not eligible: return fwd() + if len(eligible) == 1: + inner_stream = next(iter(eligible)) + else: + inner_stream = choose_contraction(body.args, eligible) + + inner_factor_ids = frozenset( + i for i, (_, fvs) in enumerate(factors) if inner_stream in fvs + ) inner_factors = [factors[i][0] for i in sorted(inner_factor_ids)] inner_stream_keys = {inner_stream} inner_deps = set().union( - *(factors[i][1] for i in f_v), fvsof(streams[v]) & stream_keys + *(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 f_v] + 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 f_v) + *(vars for i, (_, vars) in enumerate(factors) if i not in inner_factor_ids) ) # find all streams that are used in the inner factors/streams and are @@ -857,6 +850,28 @@ def reduce(self, monoid, body, streams): return result +@Operation.define +def as_float(x: int) -> float: + if isinstance(x, Term): + raise NotHandled + return float(x) + + +class PlusCastFloat(ObjectInterpretation): + @implements(Monoid.plus) + def plus(self, monoid, *args): + typs = [typeof(a) for a in args] + if any(issubclass(t, float) for t in typs) and any( + issubclass(t, int) for t in typs + ): + args = [ + as_float(a) if issubclass(t, int) else a + for (a, t) in zip(args, typs, strict=True) + ] + return monoid.plus(*args) + return fwd() + + class _ExtensibleInterpretation(UserDict, Interpretation): def extend(self, *intps: Interpretation) -> typing.Self: for intp in intps: @@ -865,10 +880,10 @@ def extend(self, *intps: Interpretation) -> typing.Self: NormalizeIntp = _ExtensibleInterpretation().extend( + ReducePartial(), MonoidOverSequence(), MonoidOverMapping(), MonoidOverCallable(), - ReduceNoStreams(), ReduceFusion(), ReduceSplit(), ReduceFactorization(), @@ -877,10 +892,8 @@ def extend(self, *intps: Interpretation) -> typing.Self: ReduceCartesianWeightedStream(), PlusEmpty(), PlusSingle(), - PlusIdentity(), PlusAssoc(), PlusDistr(), - PlusZero(), PlusConsecutiveDups(), PlusDups(), SumPlus(), @@ -890,6 +903,7 @@ def extend(self, *intps: Interpretation) -> typing.Self: ArgMinPlus(), ArgMaxPlus(), CartesianProductPlus(), + PlusCastFloat(), ) """``NormalizeIntp``applies pure-Term rewrites (associativity, distributivity, identity elimination, fusion, factorization, etc.). diff --git a/effectful/ops/syntax.py b/effectful/ops/syntax.py index 958f2fbf6..5d0b1e983 100644 --- a/effectful/ops/syntax.py +++ b/effectful/ops/syntax.py @@ -905,6 +905,14 @@ def _(x: object, other) -> bool: return x == other +@syntactic_eq.register(int | float) +def _(x: int | float, other) -> bool: + # Terms often override __eq__ + if isinstance(other, Term) or not isinstance(other, int | float): + return False + return x == other + + @_CustomSingleDispatchCallable def syntactic_hash(__dispatch: Callable[[type], Callable[[Any], int]], x) -> int: """Structural hash compatible with :func:`syntactic_eq`. diff --git a/pyproject.toml b/pyproject.toml index 054763c27..29dca1eb3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,7 +39,10 @@ Source = "https://github.com/BasisResearch/effectful" [project.optional-dependencies] torch = ["torch"] pyro = ["pyro-ppl>=1.9.1"] -jax = ["jax"] +jax = [ + "jax", + "opt_einsum" +] numpyro = [ "numpyro>=0.19", "jax<0.10" diff --git a/tests/test_handlers_jax_monoid.py b/tests/test_handlers_jax_monoid.py index 18df84018..48913d80e 100644 --- a/tests/test_handlers_jax_monoid.py +++ b/tests/test_handlers_jax_monoid.py @@ -1,78 +1,108 @@ import functools -import typing import jax import pytest from jax import random as random import effectful.handlers.jax.numpy as jnp -from effectful.handlers.jax import bind_dims, unbind_dims +from effectful.handlers.jax import bind_dims, jax_getitem, unbind_dims from effectful.handlers.jax.monoid import ( - ArrayReduce, - LogSumExp, - ProductPlusJax, - ReduceDeltaIndependent, + ARRAY_REDUCTORS, + DeltaEmpty, + ReduceArray, + ReduceArrayGather, + ReduceDeltaSimpleRange, ReduceDependentRangeMask, + ReduceSumProductContraction, delta, + einsum, ) -from effectful.handlers.jax.monoid import range as Range -from effectful.handlers.jax.scipy.special import logsumexp -from effectful.ops.monoid import ( - Max, - Min, - NormalizeIntp, - Product, - ReduceWeightedStream, - Sum, -) +from effectful.ops.monoid import NormalizeIntp, Product, Sum from effectful.ops.semantics import coproduct, handler -from effectful.ops.types import Interpretation from tests._monoid_helpers import JaxBackend MONOIDS = [ - pytest.param(Sum, jnp.sum, id="Sum"), - pytest.param(Product, jnp.prod, id="Product"), - pytest.param(Min, jnp.min, id="Min"), - pytest.param(Max, jnp.max, id="Max"), - pytest.param(LogSumExp, logsumexp, id="LogSumExp"), + pytest.param(monoid, reductor, id=monoid._name) + for (monoid, reductor) in ARRAY_REDUCTORS.items() ] +@pytest.fixture(scope="module") +def rng_key(): + return random.PRNGKey(0) + + @pytest.fixture def backend() -> JaxBackend: return JaxBackend() +@pytest.mark.parametrize("monoid,reductor", MONOIDS) +def test_reduce_array_gather(monoid, reductor, backend: JaxBackend): + (x, k) = backend.define_vars("x", "k", ret="scalar") + X = jnp.arange(3) + + lhs = monoid.reduce(x(), {x: X}) + rhs = monoid.reduce(unbind_dims(X, k), {k: range(X.shape[0])}) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceArrayGather()) + + +@pytest.mark.parametrize("monoid,reductor", MONOIDS) +def test_reduce_array_gather_dep(monoid, reductor, backend: JaxBackend): + (x, y) = backend.define_vars("x", "y", ret="scalar") + f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="stream") + g = backend.define_vars( + "g", arg_types=(backend.scalar_typ, backend.scalar_typ), ret="scalar" + ) + X = jnp.arange(3) + + lhs = monoid.reduce(g(x(), y()), {y: f(x()), x: X}) + rhs = monoid.reduce( + g(unbind_dims(X[:3], x), y()), {y: f(x()), x: range(X.shape[0])} + ) + backend.check_rewrite( + lhs=lhs, rhs=rhs, rule=coproduct(ReduceArrayGather(), ReduceDeltaSimpleRange()) + ) + + @pytest.mark.parametrize("monoid,reductor", MONOIDS) def test_reduce_array_1(monoid, reductor, backend: JaxBackend): (x, k) = backend.define_vars("x", "k", ret="scalar") - X = backend.define_vars("X", ret="stream") + X = jnp.arange(5) - lhs = monoid.reduce(x(), {x: X()}) - rhs = reductor(bind_dims(unbind_dims(X(), k), k), axis=0) - backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ArrayReduce()) + lhs = monoid.reduce(x(), {x: X}) + rhs = reductor(bind_dims(unbind_dims(X, k), k), axis=(0,)) + backend.check_rewrite( + lhs=lhs, + rhs=rhs, + rule=functools.reduce( + coproduct, # type: ignore[arg-type] + [ReduceArrayGather(), ReduceArray(), ReduceDeltaSimpleRange()], + ), + ) @pytest.mark.parametrize("monoid,reductor", MONOIDS) def test_reduce_array_2(monoid, reductor, backend: JaxBackend): (x, y, k1, k2) = backend.define_vars("x", "y", "k1", "k2", ret="scalar") - (X, Y) = backend.define_vars("X", "Y", ret="stream") + X = jnp.arange(5) + Y = jnp.arange(7) f = backend.define_vars( "f", arg_types=(backend.scalar_typ, backend.scalar_typ), ret="scalar" ) - lhs = monoid.reduce(f(x(), y()), {x: X(), y: Y()}) + lhs = monoid.reduce(f(x(), y()), {x: X, y: Y}) rhs = reductor( - bind_dims( - reductor( - bind_dims(f(unbind_dims(X(), k1), unbind_dims(Y(), k2)), k2), - axis=0, - ), - k1, + bind_dims(f(unbind_dims(X, k1), unbind_dims(Y, k2)), k1, k2), axis=(0, 1) + ) + backend.check_rewrite( + lhs=lhs, + rhs=rhs, + rule=functools.reduce( + coproduct, # type: ignore[arg-type] + [ReduceArrayGather(), ReduceArray(), ReduceDeltaSimpleRange()], ), - axis=0, ) - backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ArrayReduce()) @pytest.mark.parametrize("monoid,reductor", MONOIDS) @@ -80,58 +110,93 @@ def test_reduce_array_3(monoid, reductor, backend: JaxBackend): """Stream `y` is `g(x())` — depends on the bound element of X. The reducer must inline ``g`` along the same named dim used to unbind `x`.""" (x, y, k1, k2) = backend.define_vars("x", "y", "k1", "k2", ret="scalar") - X = backend.define_vars("X", ret="stream") + X = jnp.arange(5) f = backend.define_vars( "f", arg_types=[backend.scalar_typ, backend.scalar_typ], ret="scalar" ) g = backend.define_vars("g", arg_types=[backend.scalar_typ], ret="stream") - lhs = monoid.reduce(f(x(), y()), {x: X(), y: g(x())}) + lhs = monoid.reduce(f(x(), y()), {x: X, y: g(x())}) rhs = reductor( bind_dims( - reductor( - bind_dims( - f(unbind_dims(X(), k1), unbind_dims(g(unbind_dims(X(), k1)), k2)), - k2, - ), - axis=0, - ), - k1, + monoid.reduce(f(unbind_dims(X, x), y()), {y: g(unbind_dims(X, x))}), x + ), + axis=(0,), + ) + backend.check_rewrite( + lhs=lhs, + rhs=rhs, + rule=functools.reduce( + coproduct, # type: ignore[arg-type] + [ + ReduceArrayGather(), + ReduceArray(), + ReduceDeltaSimpleRange(), + DeltaEmpty(), + ], ), - axis=0, ) - backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ArrayReduce()) -def test_jax_weighted_reduce(backend: JaxBackend): - """Sum over a single stream with ``Product`` weights lowers to - ``jnp.sum(w(X) * body(X))`` under ``NormalizeIntp`` ∘ ``ArrayReduce``. +@pytest.mark.parametrize("monoid,reductor", MONOIDS) +def test_arange_reduce_direct_full(monoid, reductor, backend: JaxBackend): + """A full-range direct index ``A[v()]`` over ``v: arange(N)`` slices the + whole axis (``A[0:N:1]``) and reduces it -- no materialized-arange gather. + """ + (v, k) = backend.define_vars("v", "k", ret="scalar") + A = backend.define_vars("A", ret="stream") - Verifies that the desugaring rule composes cleanly with the JAX lowering - so existing handlers need no changes to support weighted streams. + lhs = monoid.reduce(jax_getitem(A(), [v()]), {v: range(7)}) + rhs = reductor( + bind_dims(jax_getitem(jax_getitem(A(), [slice(0, 7, 1)]), [k()]), k), + axis=(0,), + ) + backend.check_rewrite( + lhs=lhs, rhs=rhs, rule=coproduct(ReduceArray(), ReduceDeltaSimpleRange()) + ) - """ - (x, k) = backend.define_vars("x", "k", ret="scalar") - X = backend.define_vars("X", ret="stream") - body = backend.define_vars("body", arg_types=[backend.scalar_typ], ret="scalar") - w = backend.define_vars("w", arg_types=[backend.scalar_typ], ret="scalar") - ws = Product.weighted(X(), w) - lhs = Sum.reduce(body(x()), {x: ws}) - rhs = jnp.sum( - bind_dims(w(unbind_dims(X(), k)) * body(unbind_dims(X(), k)), k), axis=0 +@pytest.mark.parametrize("monoid,reductor", MONOIDS) +def test_arange_reduce_indirect(monoid, reductor, backend: JaxBackend): + """When the range var is used both as a direct index and as a value + (``A[v()] + v()``), the direct use slices and the indirect use materializes + the range, both aligned on the same fresh dim.""" + (v, k) = backend.define_vars("v", "k", ret="scalar") + A = jnp.arange(10) + + lhs = monoid.reduce(jax_getitem(A, [v()]) + v(), {v: range(5)}) + rhs = reductor( + bind_dims( + jax_getitem(jax_getitem(A, [slice(0, 5, 1)]), [k()]) + + unbind_dims(jnp.arange(5), k), + k, + ), + axis=(0,), ) backend.check_rewrite( - lhs=lhs, - rhs=rhs, - rule=functools.reduce( - coproduct, - typing.cast( - list[Interpretation], - [ReduceWeightedStream(), ArrayReduce(), ProductPlusJax()], - ), + lhs=lhs, rhs=rhs, rule=coproduct(ReduceArray(), ReduceDeltaSimpleRange()) + ) + + +@pytest.mark.parametrize("monoid,reductor", MONOIDS) +def test_arange_reduce_two_streams(monoid, reductor, backend: JaxBackend): + """Two arange streams indexing a 2-D array slice both axes and reduce over + both at once.""" + (u, w, k1, k2) = backend.define_vars("u", "w", "k1", "k2", ret="scalar") + A = jnp.arange(8 * 9).reshape((8, 9)) + + lhs = monoid.reduce(jax_getitem(A, [u(), w()]), {u: range(4), w: range(5)}) + rhs = reductor( + bind_dims( + jax_getitem(jax_getitem(A, [slice(0, 4, 1), slice(0, 5, 1)]), [k1(), k2()]), + k1, + k2, ), + axis=(0, 1), + ) + backend.check_rewrite( + lhs=lhs, rhs=rhs, rule=coproduct(ReduceArray(), ReduceDeltaSimpleRange()) ) @@ -152,15 +217,26 @@ def test_reduce_delta_empty(monoid, reductor, backend: JaxBackend): lhs = monoid.reduce(delta((), x()), {x: X()}) rhs = monoid.reduce(x(), {x: X()}) - backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceDeltaIndependent()) + backend.check_rewrite( + lhs=lhs, rhs=rhs, rule=coproduct(ReduceDeltaSimpleRange(), DeltaEmpty()) + ) + + +@pytest.mark.parametrize("monoid,reductor", MONOIDS) +def test_reduce_delta_empty_arange(monoid, reductor, backend: JaxBackend): + x = backend.define_vars("x", ret="scalar") + f = backend.define_vars("f", arg_types=[backend.scalar_typ], ret="scalar") + + lhs = monoid.reduce(delta((x(),), f(x())), {x: range(0)}) + rhs = bind_dims(f(unbind_dims(jnp.array([]), x)), x) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceDeltaSimpleRange()) @pytest.mark.parametrize("monoid,reductor", MONOIDS) def test_reduce_delta_independent_one(monoid, reductor, backend: JaxBackend): """One R1 step: peel the final preserved index off a delta. - reduce(M, {y: Y()}, delta((y(),), f(y()))) - ≡ reduce(M, {}, delta((), bind_dims(f(unbind_dims(Y(), k)), k))) + reduce(M, {y: Y()}, delta((y(),), f(y()))) ≡ bind_dims(f(unbind_dims(Y(), k)), k) """ (y, k) = backend.define_vars("y", "k", ret="scalar") f = backend.define_vars("f", arg_types=[backend.scalar_typ], ret="scalar") @@ -168,9 +244,9 @@ def test_reduce_delta_independent_one(monoid, reductor, backend: JaxBackend): # We use a concrete range here instead of an abstract one, because # unbind_dims is undefined on empty arrays (and the rewrite produces a # different rhs in this case) - lhs = monoid.reduce(delta((y(),), f(y())), {y: Range(3)}) - rhs = monoid.reduce(bind_dims(f(unbind_dims(jnp.arange(3), k)), k), {}) - backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceDeltaIndependent()) + lhs = monoid.reduce(delta((y(),), f(y())), {y: range(3)}) + rhs = bind_dims(f(unbind_dims(jnp.arange(3), k)), k) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceDeltaSimpleRange()) @pytest.mark.parametrize("monoid,reductor", MONOIDS) @@ -188,17 +264,34 @@ def test_reduce_delta_independent_preserves_others( "f", arg_types=[backend.scalar_typ, backend.scalar_typ], ret="scalar" ) - lhs = monoid.reduce(delta((x(), y()), f(x(), y())), {x: Range(2), y: Range(3)}) - rhs = monoid.reduce( - bind_dims( - bind_dims( - f(unbind_dims(jnp.arange(2), x), unbind_dims(jnp.arange(3), k)), k - ), - x, + lhs = monoid.reduce(delta((x(), y()), f(x(), y())), {x: range(2), y: range(3)}) + rhs = bind_dims( + bind_dims(f(unbind_dims(jnp.arange(2), x), unbind_dims(jnp.arange(3), k)), k), x + ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceDeltaSimpleRange()) + + +@pytest.mark.parametrize("monoid,reductor", MONOIDS) +def test_reduce_delta_simple_dep(monoid, reductor, backend: JaxBackend): + (x, y) = backend.define_vars("x", "y", ret="scalar") + X = jnp.arange(3) + + lhs = monoid.reduce( + delta((x(),), unbind_dims(X, x) + y()), + {x: range(3), y: jnp.stack([x(), x() + 1])}, + ) + rhs = bind_dims( + monoid.reduce( + delta((), unbind_dims(X, x) + y()), + { + y: jnp.stack( + [unbind_dims(jnp.arange(3), x), unbind_dims(jnp.arange(3), x) + 1] + ) + }, ), - {}, + x, ) - backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceDeltaIndependent()) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceDeltaSimpleRange()) @pytest.mark.parametrize("monoid,reductor", MONOIDS) @@ -217,10 +310,9 @@ def test_reduce_dependent_range_mask(monoid, reductor, backend: JaxBackend): body = f(u(), v()) - lhs = monoid.reduce(body, {u: Range(0, N, 1), v: Range(0, u(), 1)}) + lhs = monoid.reduce(body, {u: range(N), v: jnp.arange(u())}) rhs = monoid.reduce( - jnp.where(v() < u(), body, monoid.identity), - {u: Range(0, N, 1), v: Range(0, N, 1)}, + jnp.where(v() < u(), body, monoid.identity), {u: range(N), v: range(N)} ) backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceDependentRangeMask()) @@ -243,14 +335,44 @@ def test_reduce_dependent_range_mask_delta_body(monoid, reductor, backend: JaxBa weight = f(u(), v()) idx = (u(), v()) - lhs = monoid.reduce(delta(idx, weight), {u: Range(0, N, 1), v: Range(0, u(), 1)}) + lhs = monoid.reduce(delta(idx, weight), {u: range(N), v: jnp.arange(u())}) rhs = monoid.reduce( delta(idx, jnp.where(v() < u(), weight, monoid.identity)), - {u: Range(0, N, 1), v: Range(0, N, 1)}, + {u: range(N), v: range(N)}, ) backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceDependentRangeMask()) +def test_reduce_contraction_single(backend: JaxBackend): + i = backend.define_vars("i", ret="scalar") + (A, B) = backend.define_vars( + "A", "B", arg_types=(backend.scalar_typ,), ret="scalar" + ) + + lhs = Sum.reduce(Product.plus(A(i()), B(i())), {i: range(5)}) + rhs = jnp.einsum( + "a...,a...->...", + Sum.reduce(delta((i(),), A(i())), {i: range(5)}), + Sum.reduce(delta((i(),), B(i())), {i: range(5)}), + ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceSumProductContraction()) + + +def test_reduce_contraction_double(backend: JaxBackend): + i, j = backend.define_vars("i", "j", ret="scalar") + (A, B) = backend.define_vars( + "A", "B", arg_types=(backend.scalar_typ, backend.scalar_typ), ret="scalar" + ) + + lhs = Sum.reduce(Product.plus(A(i(), j()), B(i(), j())), {i: range(5), j: range(7)}) + rhs = jnp.einsum( + "ab...,ab...->...", + Sum.reduce(delta((i(), j()), A(i(), j())), {i: range(5), j: range(7)}), + Sum.reduce(delta((i(), j()), B(i(), j())), {i: range(5), j: range(7)}), + ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceSumProductContraction()) + + def test_reduce_matmul(backend: JaxBackend): key = jax.random.PRNGKey(0) # Define dimensions @@ -264,8 +386,135 @@ def test_reduce_matmul(backend: JaxBackend): with handler(NormalizeIntp): actual = Sum.reduce( delta((b(), i(), k()), unbind_dims(X, b, i, j) * unbind_dims(Y, b, j, k)), - {b: Range(B), i: Range(I), j: Range(J), k: Range(K)}, + {b: range(B), i: range(I), j: range(J), k: range(K)}, ) expected = jnp.einsum("bij,bjk->bik", X, Y) assert jnp.allclose(actual, expected) + + +EINSUM_CASES = [ + pytest.param("ij,jk->ik", {"i": 64, "j": 64, "k": 64}, id="matmul"), + pytest.param( + "bij,bjk->bik", + {"b": 16, "i": 32, "j": 32, "k": 32}, + id="batched_matmul", + ), + pytest.param( + "a,abi,bcij,cdij->ij", + {"a": 4, "b": 4, "c": 4, "d": 4, "i": 8, "j": 8}, + id="mixed_rank", + ), + # ───────────────────────── single-operand reshuffles ───────────────────── + # No contraction across operands — these stress the diagonal/transpose/sum + # rewrites rather than any pairwise product ordering. + pytest.param("ij->ji", {"i": 256, "j": 256}, id="transpose"), + pytest.param("ijk->", {"i": 96, "j": 96, "k": 96}, id="full_reduce"), + pytest.param("ijk->k", {"i": 96, "j": 96, "k": 96}, id="partial_reduce"), + # Repeated index *within* one operand — exercises the implicit-diagonal path + # in ReduceDeltaSimpleRange (no explicit jnp.diagonal step). + pytest.param("ii->", {"i": 1024}, id="trace"), + pytest.param("ii->i", {"i": 1024}, id="diagonal"), + pytest.param("bii->b", {"b": 256, "i": 128}, id="batched_trace"), + pytest.param("iij->ij", {"i": 128, "j": 128}, id="diagonal_keep"), + # ───────────────────────── no-shared-index blowups ─────────────────────── + # Output is the full outer product — nothing contracts, so the result tensor + # is as large as the dense intermediate. Pure broadcast cost. + pytest.param("i,j->ij", {"i": 1024, "j": 1024}, id="outer_product"), + pytest.param("ij,kl->ijkl", {"i": 32, "j": 32, "k": 32, "l": 32}, id="outer_4d"), + # Element-wise: every index shared, none contracted. + pytest.param("ij,ij->ij", {"i": 512, "j": 512}, id="hadamard"), + # ───────────────────────── ordering-sensitive products ─────────────────── + # Skewed matrix chain: contracting middle-first (b,d small) is orders of + # magnitude cheaper than the left-to-right order, which materializes a big + # a×c intermediate. The classic "matrix chain order matters" case. + pytest.param( + "ab,bc,cd->ad", {"a": 256, "b": 2, "c": 256, "d": 2}, id="skewed_chain" + ), + pytest.param( + "ab,bc,cd,de->ae", + {"a": 50, "b": 40, "c": 30, "d": 20, "e": 10}, + id="chain_4", + ), + pytest.param( + "ab,bc,cd,de,ef->af", + {"a": 12, "b": 11, "c": 10, "d": 9, "e": 8, "f": 7}, + id="chain_5", + ), + # ───────────────────────── tensor-network shapes ───────────────────────── + # Cyclic / hyperedge contractions with no tree decomposition into matmuls; + # every operand shares indices with two others. + pytest.param("ij,jk,ki->", {"i": 64, "j": 64, "k": 64}, id="trace_of_product"), + pytest.param("ij,jk,ik->", {"i": 48, "j": 48, "k": 48}, id="triangle"), + pytest.param("ijk,jl,kl->il", {"i": 24, "j": 24, "k": 24, "l": 24}, id="hyperedge"), + # Star: many operands share one contracted index, fanning into a large + # outer-product output. + pytest.param( + "ai,bi,ci,di->abcd", + {"a": 8, "b": 8, "c": 8, "d": 8, "i": 32}, + id="star_contraction", + ), + # Bilinear / quadratic form over a batch (attention-score flavored). + pytest.param("bi,ij,bj->b", {"b": 128, "i": 64, "j": 64}, id="bilinear"), + # Batched matrix chain — batch axis rides through three contractions. + pytest.param( + "bij,bjk,bkl->bil", + {"b": 16, "i": 24, "j": 24, "k": 24, "l": 24}, + id="batched_chain", + ), + # Multi-index contraction surface: a whole axis-group (c) contracts at once. + pytest.param( + "abc,cde->abde", + {"a": 12, "b": 12, "c": 12, "d": 12, "e": 12}, + id="tensor_contraction", + ), + # Leading scalar factor plus an element-wise reduce — checks that the + # rank-0 operand threads through without spawning a degenerate axis. + pytest.param(",ij,ij->", {"i": 256, "j": 256}, id="scalar_scaled_reduce"), +] + + +def _make_operands(spec: str, sizes: dict[str, int], key: jax.Array) -> list[jax.Array]: + in_part = spec.split("->")[0] + in_specs = in_part.split(",") + keys = random.split(key, len(in_specs)) + return [ + random.normal(k, tuple(sizes[c] for c in s) if s else ()) + for k, s in zip(keys, in_specs, strict=True) + ] + + +@pytest.mark.parametrize( + "impl", [pytest.param(jnp.einsum, id="jax"), pytest.param(einsum, id="effectful")] +) +@pytest.mark.parametrize("spec,sizes", EINSUM_CASES) +@pytest.mark.benchmark(warmup=True, warmup_iterations=1) +def test_einsum_bench(benchmark, impl, spec, sizes, rng_key): + """Time one ``(spec, impl)`` pair. Group by ``spec`` to compare ``jnp`` + against ``effectful`` for the same subscript pattern (see module docstring). + """ + operands = _make_operands(spec, sizes, rng_key) + + @jax.jit + def f(*operands): + return impl(spec, *operands) + + @benchmark + def _run(): + return f(*operands).block_until_ready() + + +@pytest.mark.parametrize("spec,sizes", EINSUM_CASES) +def test_einsum_matches_jnp(spec: str, sizes, rng_key): + """``einsum`` returns the same result as ``jnp.einsum`` for every spec + in ``EINSUM_EXAMPLES``. + """ + operands = _make_operands(spec, sizes, rng_key) + actual = einsum(spec, *operands) + expected = jnp.einsum(spec, *operands) + assert actual.shape == expected.shape, ( + f"shape mismatch for {spec!r}: got {actual.shape}, expected {expected.shape}" + ) + assert jnp.allclose(actual, expected, atol=1e-4, rtol=1e-4), ( + f"value mismatch for {spec!r}" + ) diff --git a/tests/test_ops_monoid.py b/tests/test_ops_monoid.py index 4d243ca14..484a8c4e8 100644 --- a/tests/test_ops_monoid.py +++ b/tests/test_ops_monoid.py @@ -21,15 +21,13 @@ PlusDistr, PlusDups, PlusEmpty, - PlusIdentity, PlusSingle, - PlusZero, Product, ReduceCartesianWeightedStream, ReduceDistributeCartesianProduct, ReduceFactorization, ReduceFusion, - ReduceNoStreams, + ReducePartial, ReduceSplit, ReduceWeightedStream, Sum, @@ -172,7 +170,7 @@ def test_plus_identity_right(monoid, backend: Backend): lhs = monoid.plus(x(), monoid.identity) rhs = monoid.plus(x()) - backend.check_rewrite(lhs=lhs, rhs=rhs, rule=PlusIdentity()) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule={}) @pytest.mark.parametrize("monoid", ALL_MONOIDS) @@ -182,7 +180,7 @@ def test_plus_identity_left(monoid, backend: Backend): lhs = monoid.plus(monoid.identity, x()) rhs = monoid.plus(x()) - backend.check_rewrite(lhs=lhs, rhs=rhs, rule=PlusIdentity()) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule={}) @pytest.mark.parametrize("monoid", ALL_MONOIDS) @@ -240,10 +238,10 @@ def test_plus_distributes(backend: Backend): def test_plus_distributes_constant(backend: Backend): - a, b, c, d = backend.define_vars("a", "b", "c", "d", ret="scalar") - lhs = Product.plus(Sum.plus(a(), b()), Sum.plus(c(), d()), 5) + a, b, c, d, e = backend.define_vars("a", "b", "c", "d", "e", ret="scalar") + lhs = Product.plus(Sum.plus(a(), b()), Sum.plus(c(), d()), e()) rhs = Product.plus( - 5, + e(), Sum.plus( Product.plus(a(), c()), Product.plus(a(), d()), @@ -314,16 +312,16 @@ def test_plus_zero(monoid, backend: Backend): lhs_right = monoid.plus(a(), monoid.zero) lhs_left = monoid.plus(monoid.zero, a()) rhs = monoid.zero - backend.check_rewrite(lhs=lhs_right, rhs=rhs, rule=PlusZero()) - backend.check_rewrite(lhs=lhs_left, rhs=rhs, rule=PlusZero()) + backend.check_rewrite(lhs=lhs_right, rhs=rhs, rule={}) + backend.check_rewrite(lhs=lhs_left, rhs=rhs, rule={}) @pytest.mark.parametrize("monoid", ALL_MONOIDS) def test_partial_1(monoid, backend: Backend): x = backend.define_vars("x", ret="scalar") lhs = monoid.reduce(x(), {x: []}) - rhs = monoid.identity - backend.check_rewrite(lhs=lhs, rhs=rhs, rule={}) + rhs = monoid.plus() + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReducePartial()) @pytest.mark.parametrize("monoid", ALL_MONOIDS) @@ -332,8 +330,8 @@ def test_partial_2(monoid, backend: Backend): Y = backend.define_vars("Y", ret="stream") lhs = monoid.reduce(x(), {y: Y(), x: []}) - rhs = monoid.identity - backend.check_rewrite(lhs=lhs, rhs=rhs, rule={}) + rhs = monoid.plus() + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReducePartial()) @pytest.mark.parametrize("monoid", ALL_MONOIDS) @@ -343,7 +341,7 @@ def test_partial_3(monoid, backend: Backend): lhs = monoid.reduce(x(), {y: Y(), x: [a(), b()]}) rhs = monoid.plus(monoid.reduce(a(), {y: Y()}), monoid.reduce(b(), {y: Y()})) - backend.check_rewrite(lhs=lhs, rhs=rhs, rule={}) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReducePartial()) @pytest.mark.parametrize("monoid", ALL_MONOIDS) @@ -353,7 +351,7 @@ def test_partial_4(monoid, backend: Backend): lhs = monoid.reduce(x(), {y: f(x()), x: [a(), b()]}) rhs = monoid.plus(monoid.reduce(a(), {y: f(a())}), monoid.reduce(b(), {y: f(b())})) - backend.check_rewrite(lhs=lhs, rhs=rhs, rule={}) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReducePartial()) @pytest.mark.parametrize("monoid", ALL_MONOIDS) @@ -401,7 +399,7 @@ def test_reduce_no_streams(monoid, backend: Backend): lhs = monoid.reduce(a(), {}) rhs = monoid.identity - backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceNoStreams()) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReducePartial()) @pytest.mark.parametrize("monoid", ALL_MONOIDS) @@ -482,15 +480,15 @@ def test_reduce_independent_3_negative(backend: Backend): def test_reduce_independent_4(backend: Backend): - a, b, c = backend.define_vars("a", "b", "c", ret="scalar") + a, b, c, d = backend.define_vars("a", "b", "c", "d", ret="scalar") A, B, C = backend.define_vars("A", "B", "C", ret="stream") f = backend.define_vars( "f", arg_types=(backend.scalar_typ, backend.scalar_typ), ret="scalar" ) - lhs = Sum.reduce(Product.plus(a(), b(), f(b(), c()), 7), {a: A(), b: B(), c: C()}) + lhs = Sum.reduce(Product.plus(a(), b(), f(b(), c()), d()), {a: A(), b: B(), c: C()}) rhs = Product.plus( - 7, + d(), Sum.reduce(Product.plus(a()), {a: A()}), Sum.reduce( Product.plus(b(), Sum.reduce(Product.plus(f(b(), c())), {c: C()})), From 08d4a0caaa3c27bd5c48d1864532764c20c81bcc Mon Sep 17 00:00:00 2001 From: eb8680 Date: Thu, 11 Jun 2026 14:35:51 -0400 Subject: [PATCH 09/16] Break ReduceArrayGather rule into 2 steps (#681) --- effectful/handlers/jax/monoid.py | 29 ++++++++++------- effectful/ops/monoid.py | 52 +++++++++++++++++++++++++++++++ tests/test_handlers_jax_monoid.py | 48 +++++++++++++++++++++++++--- tests/test_ops_monoid.py | 28 +++++++++++++++++ 4 files changed, 140 insertions(+), 17 deletions(-) diff --git a/effectful/handlers/jax/monoid.py b/effectful/handlers/jax/monoid.py index fbd845d8b..96b708e89 100644 --- a/effectful/handlers/jax/monoid.py +++ b/effectful/handlers/jax/monoid.py @@ -156,7 +156,17 @@ def plus(self, *args): class ReduceArrayGather(ObjectInterpretation): - """M.reduce(body, {k: a} ∪ S) ≡ M.reduce(body[k := a[k']], {k': range(a.shape[0])} ∪ S)""" + """Split an array-valued stream into an index range and a length-1 stream: + + M.reduce(body, {k: a} ∪ S) ≡ M.reduce(body, {i: range(a.shape[0]), k: (a[i()],)} ∪ S) + + where ``i`` is fresh and ``a[i()] = unbind_dims(a, i)``. The length-1 stream + ``{k: (a[i()],)}`` is then eliminated by + :class:`~effectful.ops.monoid.EliminateSingletonStreams`, which substitutes + ``k := a[i()]`` into the body and the remaining streams. Together the two + steps perform the gather + ``M.reduce(body[k := a[i()]], {i: range(a.shape[0])} ∪ S)``. + """ @implements(Monoid.reduce) def reduce(self, monoid, body, streams): @@ -169,26 +179,21 @@ def reduce(self, monoid, body, streams): body_fvs = fvsof(body) stream_keys = set(streams) - body_subst = {} - streams_subst = {} - range_streams = {} + new_streams: dict = {} progress = False for k, v in streams.items(): if is_eager_array(v) and k in body_fvs and not (fvsof(v) & stream_keys): - kk = Operation.define(k) - body_subst[k] = deffn(unbind_dims(v, kk)) - streams_subst[k] = kk - range_streams[kk] = range(v.shape[0]) + index = Operation.define(k) + new_streams[index] = range(v.shape[0]) + new_streams[k] = (unbind_dims(v, index),) progress = True else: - range_streams[k] = v + new_streams[k] = v if not progress: return fwd() - subst_body = handler(body_subst)(evaluate)(body) - subst_streams = handler(streams_subst)(evaluate)(range_streams) - return monoid.reduce(subst_body, subst_streams) + return monoid.reduce(body, new_streams) class Reductor(Protocol): diff --git a/effectful/ops/monoid.py b/effectful/ops/monoid.py index 42180e912..067f29e55 100644 --- a/effectful/ops/monoid.py +++ b/effectful/ops/monoid.py @@ -872,6 +872,57 @@ def plus(self, monoid, *args): return fwd() +class EliminateSingletonStreams(ObjectInterpretation): + """Eliminate a length-1 stream by substituting its sole element. + + reduce(M, body, {k: (v,)} ∪ S) = reduce(M, body[k := v], S[k := v]) + + Fires only when the sole element ``v`` is a :class:`Term`, i.e. a *symbolic* + singleton. This is exactly the form ``ReduceArrayGather`` produces (a gather + ``(a[i()],)``) and, more generally, every dependent singleton that + :class:`ReducePartial` cannot peel -- a non-outermost stream whose element + references another stream var. Concrete enumerated streams (``[0]``, + ``range(1)``) and monoid sentinels (``CartesianProduct.identity == [()]``) + have non-``Term`` elements and are left to ``ReducePartial`` / the + per-monoid rules. + + Unlike ``ReducePartial``, this peels the stream wherever it sits in the loop + nest and substitutes symbolically rather than unrolling, leaving a + vectorized index range (e.g. the gather's range) intact instead of + materializing it. + """ + + @implements(Monoid.reduce) + def reduce(self, monoid, body, streams): + # Eliminate *all* symbolic length-1 streams in one pass via a + # simultaneous substitution. Doing them together (rather than one per + # invocation) keeps an interleaving reduction rule -- e.g. + # ``ReduceArray`` consuming a now-live index range -- from firing + # between eliminations, so sibling index ranges stay together and fuse + # into a single reduction. + singletons = { + k: vs[0] + for k, vs in streams.items() + if not isinstance(vs, Term) + and isinstance(vs, collections.abc.Sequence) + and len(vs) == 1 + and isinstance(vs[0], Term) + } + if not singletons: + return fwd() + + subs = {k: deffn(v) for k, v in singletons.items()} + new_body = handler(subs)(evaluate)(body) + new_streams = { + kk: handler(subs)(evaluate)(vv) + for kk, vv in streams.items() + if kk not in singletons + } + # reduce over no streams is a single (empty) assignment, i.e. the body + # itself -- not the monoid identity. + return monoid.reduce(new_body, new_streams) if new_streams else new_body + + class _ExtensibleInterpretation(UserDict, Interpretation): def extend(self, *intps: Interpretation) -> typing.Self: for intp in intps: @@ -881,6 +932,7 @@ def extend(self, *intps: Interpretation) -> typing.Self: NormalizeIntp = _ExtensibleInterpretation().extend( ReducePartial(), + EliminateSingletonStreams(), MonoidOverSequence(), MonoidOverMapping(), MonoidOverCallable(), diff --git a/tests/test_handlers_jax_monoid.py b/tests/test_handlers_jax_monoid.py index 48913d80e..e410a3e69 100644 --- a/tests/test_handlers_jax_monoid.py +++ b/tests/test_handlers_jax_monoid.py @@ -17,7 +17,12 @@ delta, einsum, ) -from effectful.ops.monoid import NormalizeIntp, Product, Sum +from effectful.ops.monoid import ( + EliminateSingletonStreams, + NormalizeIntp, + Product, + Sum, +) from effectful.ops.semantics import coproduct, handler from tests._monoid_helpers import JaxBackend @@ -44,6 +49,23 @@ def test_reduce_array_gather(monoid, reductor, backend: JaxBackend): lhs = monoid.reduce(x(), {x: X}) rhs = monoid.reduce(unbind_dims(X, k), {k: range(X.shape[0])}) + backend.check_rewrite( + lhs=lhs, + rhs=rhs, + rule=coproduct(ReduceArrayGather(), EliminateSingletonStreams()), + ) + + +@pytest.mark.parametrize("monoid,reductor", MONOIDS) +def test_reduce_array_gather_step1(monoid, reductor, backend: JaxBackend): + """Step 1 alone: an array stream becomes an index range plus a length-1 + stream holding the gathered element. ``ReduceArrayGather`` does not perform + the gather substitution itself -- that is ``EliminateSingletonStreams``.""" + (x, k) = backend.define_vars("x", "k", ret="scalar") + X = jnp.arange(3) + + lhs = monoid.reduce(x(), {x: X}) + rhs = monoid.reduce(x(), {k: range(X.shape[0]), x: (unbind_dims(X, k),)}) backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceArrayGather()) @@ -56,12 +78,17 @@ def test_reduce_array_gather_dep(monoid, reductor, backend: JaxBackend): ) X = jnp.arange(3) + # The dependent stream ``y: f(x())`` gets the *gathered element* X[x] + # substituted for x -- i.e. ``f(X[x])`` -- not the bare index. lhs = monoid.reduce(g(x(), y()), {y: f(x()), x: X}) rhs = monoid.reduce( - g(unbind_dims(X[:3], x), y()), {y: f(x()), x: range(X.shape[0])} + g(unbind_dims(X, x), y()), + {y: f(unbind_dims(X, x)), x: range(X.shape[0])}, ) backend.check_rewrite( - lhs=lhs, rhs=rhs, rule=coproduct(ReduceArrayGather(), ReduceDeltaSimpleRange()) + lhs=lhs, + rhs=rhs, + rule=coproduct(ReduceArrayGather(), EliminateSingletonStreams()), ) @@ -77,7 +104,12 @@ def test_reduce_array_1(monoid, reductor, backend: JaxBackend): rhs=rhs, rule=functools.reduce( coproduct, # type: ignore[arg-type] - [ReduceArrayGather(), ReduceArray(), ReduceDeltaSimpleRange()], + [ + ReduceArrayGather(), + EliminateSingletonStreams(), + ReduceArray(), + ReduceDeltaSimpleRange(), + ], ), ) @@ -100,7 +132,12 @@ def test_reduce_array_2(monoid, reductor, backend: JaxBackend): rhs=rhs, rule=functools.reduce( coproduct, # type: ignore[arg-type] - [ReduceArrayGather(), ReduceArray(), ReduceDeltaSimpleRange()], + [ + ReduceArrayGather(), + EliminateSingletonStreams(), + ReduceArray(), + ReduceDeltaSimpleRange(), + ], ), ) @@ -131,6 +168,7 @@ def test_reduce_array_3(monoid, reductor, backend: JaxBackend): coproduct, # type: ignore[arg-type] [ ReduceArrayGather(), + EliminateSingletonStreams(), ReduceArray(), ReduceDeltaSimpleRange(), DeltaEmpty(), diff --git a/tests/test_ops_monoid.py b/tests/test_ops_monoid.py index 484a8c4e8..9976fd6dc 100644 --- a/tests/test_ops_monoid.py +++ b/tests/test_ops_monoid.py @@ -10,6 +10,7 @@ import effectful.handlers.jax.numpy as jnp from effectful.ops.monoid import ( CartesianProduct, + EliminateSingletonStreams, Max, Min, Monoid, @@ -354,6 +355,33 @@ def test_partial_4(monoid, backend: Backend): backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReducePartial()) +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_eliminate_singleton_into_sibling(monoid, backend: Backend): + """A length-1 stream substitutes its element into the body *and* into a + sibling stream's definition, then drops out of the nest.""" + x, y, a = backend.define_vars("x", "y", "a", ret="scalar") + f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="stream") + g = backend.define_vars( + "g", arg_types=(backend.scalar_typ, backend.scalar_typ), ret="scalar" + ) + + lhs = monoid.reduce(g(x(), y()), {x: (a(),), y: f(x())}) + rhs = monoid.reduce(g(a(), y()), {y: f(a())}) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=EliminateSingletonStreams()) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_eliminate_singleton_only_stream(monoid, backend: Backend): + """When the length-1 stream is the only stream, reducing over the now-empty + nest yields the substituted body itself (not the monoid identity).""" + x, a = backend.define_vars("x", "a", ret="scalar") + f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="scalar") + + lhs = monoid.reduce(f(x()), {x: (a(),)}) + rhs = f(a()) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=EliminateSingletonStreams()) + + @pytest.mark.parametrize("monoid", ALL_MONOIDS) def test_reduce_body_sequence(monoid, backend: Backend): x = backend.define_vars("x", ret="scalar") From ca66476f282c3164e6aaa913c2abb99863a0dd78 Mon Sep 17 00:00:00 2001 From: eb8680 Date: Thu, 11 Jun 2026 14:36:23 -0400 Subject: [PATCH 10/16] Remove dead code (#683) --- effectful/ops/types.py | 67 ------------------------------------------ 1 file changed, 67 deletions(-) diff --git a/effectful/ops/types.py b/effectful/ops/types.py index bd87d1abf..018932eb6 100644 --- a/effectful/ops/types.py +++ b/effectful/ops/types.py @@ -42,59 +42,6 @@ def __call__(self, *args: P.args, **kwargs: P.kwargs) -> T: return self.func(self.dispatch, *args, **kwargs) -class _CustomSingleDispatchMethod[**P, **Q, S, T]: - """Method analog of :class:`_CustomSingleDispatchCallable`. - - The wrapped function has signature ``(self, dispatch, *args, **kwargs)``, - where ``dispatch`` is :meth:`functools.singledispatch.dispatch`. As a - descriptor, it binds ``self`` on attribute access, so callers invoke it - as ``instance.method(*args, **kwargs)``. - """ - - def __init__( - self, - func: Callable[Concatenate[Any, Callable[[type], Callable[Q, S]], P], T], - ): - self.func = func - self._registry = functools.singledispatch(func) - self.__signature__ = inspect.signature( - functools.partial(func, None, None) # type: ignore[arg-type] - ) - functools.update_wrapper(self, func) # type: ignore[arg-type] - - @property - def dispatch(self): - return self._registry.dispatch - - @property - def register(self): - return self._registry.register - - def __get__(self, instance, owner=None): - if instance is None: - return self - return _BoundCustomSingleDispatchMethod(self, instance) - - -class _BoundCustomSingleDispatchMethod: - __slots__ = ("_method", "_instance") - - def __init__(self, method: _CustomSingleDispatchMethod, instance: Any): - self._method = method - self._instance = instance - - @property - def dispatch(self): - return self._method.dispatch - - @property - def register(self): - return self._method.register - - def __call__(self, *args, **kwargs): - return self._method.func(self._instance, self._method.dispatch, *args, **kwargs) - - class _ClassMethodOpDescriptor(classmethod): def __init__(self, define, *args, **kwargs): super().__init__(*args, **kwargs) @@ -420,20 +367,6 @@ def func(*args, **kwargs): op.register = default._registry.register # type: ignore[attr-defined] return op - @define.register(_CustomSingleDispatchMethod) - @classmethod - def _define_customsingledispatchmethod( - cls, default: _CustomSingleDispatchMethod, **kwargs - ): - @functools.wraps(default.func) - def _wrapper(obj, *args, **kwargs): - return default.__get__(obj)(*args, **kwargs) - - op = cls.define(_wrapper, **kwargs) - op.register = default.register # type: ignore[attr-defined] - op.dispatch = default.dispatch # type: ignore[attr-defined] - return op - @typing.final def __default_rule__(self, *args: Q.args, **kwargs: Q.kwargs) -> "Expr[V]": """The default rule is used when the operation is not handled. From bd6d439602141bb4cf959593750fc78c27e7e076 Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Mon, 20 Jul 2026 11:26:57 -0400 Subject: [PATCH 11/16] Implement plated einsum (#692) * wip * wip * drop syntactic tests * wip * wip * wip * wip * wip * add cartesian product tests * wip * wip * revise reducesplit to leave shared streams * wip * allow ReduceEqualityMaskRange to look through plus * add mask hoisting * wip * wip * wip * unsupplied parameters stay bound in deffn * wip * wip * wip * allow factorization over masks * don't do leave-one-out for factors with no output dims * plated einsum tests pass * format * fix some tests * fix tests * more fixes and documentation * passing plated einsum tests * format * simplify and generalize plusdistr * normalize both order and duplicates in plus * drop unused test * push masks instead of hoisting * stop emitting extra masks * wip * add missing code * wip * revert to simple ReduceSplit, handle plus bodies in cprod elim * replace ReduceFactorization with new combined Factor * wip * give delta mapping semantics * fix tests * fix tests * add ReduceEqualityMaskRange tests * wip * wip * refactor einsum term generation * wip * wip * use where as a hoistable conditional primitive * drop binddimswhere * fix tests and clean up * avoid expensive traversal when looking up signatures * fix tests * remove ReduceCartesianWeightedStream cartesian products changed type, making these rules invalid * wip * wip * fix tests remove outdated * lint * drop test * fix notebook * drop 3.14 in CI * drop 3.14 * wip * restrict litellm * reset ci scripts * reset * drop _ArrayTerm * introduce ite op and move where handlers to ops/monoid.py * wip * remove SplitDisjointProduct * wip * move ReduceDependentRangeMask to ops/monoid.py * move ContractLongestArrayStream to ops/monoid.py * start generalizing ReduceDistributeCartesianProduct * wip * drop unused SumOfProductsIntp * wip * drop unused * simplify * more work * lint * update comment * add tests to ReduceUnfactor * drop unnecessary check * fix bug * fix tests * format * lint * drop unused code * fix tests * replace Union.delta with a dict building op * drop comment * replace Union.delta with as_dict * lint * revert * fix flipped mask * require simple ranges in scan rule * clean up collection types * enforce simple range requirement in dependent range elim * simplify typing in _EinsumBuilder * lint * add helper function _conjuncts * extract jax-independent behavior * fix test --- docs/source/introduction.ipynb | 18 +- effectful/handlers/jax/_handlers.py | 39 +- effectful/handlers/jax/_terms.py | 6 +- effectful/handlers/jax/lax/__init__.py | 19 + effectful/handlers/jax/monoid.py | 735 ++++++++---- effectful/handlers/jax/numpy/__init__.py | 17 + effectful/ops/monoid.py | 1361 ++++++++++++++++++---- effectful/ops/syntax.py | 209 +++- effectful/ops/types.py | 15 +- pyproject.toml | 6 +- tests/_monoid_helpers.py | 30 +- tests/test_handlers_jax.py | 5 +- tests/test_handlers_jax_monoid.py | 272 ++--- tests/test_ops_monoid.py | 960 ++++++++++++--- tests/test_ops_semantics.py | 5 +- 15 files changed, 2816 insertions(+), 881 deletions(-) create mode 100644 effectful/handlers/jax/lax/__init__.py diff --git a/docs/source/introduction.ipynb b/docs/source/introduction.ipynb index 400e5104d..93ed18e13 100644 --- a/docs/source/introduction.ipynb +++ b/docs/source/introduction.ipynb @@ -21,7 +21,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 1, "id": "5278fd54", "metadata": {}, "outputs": [], @@ -29,7 +29,7 @@ "import functools\n", "\n", "from effectful.ops.semantics import coproduct, fwd, handler\n", - "from effectful.ops.syntax import defdata, defop\n", + "from effectful.ops.syntax import defdata, defop, syntactic_eq\n", "from effectful.ops.types import NotHandled, Operation, Term\n", "\n", "\n", @@ -65,7 +65,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 2, "id": "3c575e02", "metadata": { "lines_to_next_cell": 2 @@ -78,7 +78,7 @@ "\n", "\n", "def test_adding_two_numbers():\n", - " assert add(1, 3) == add(2, 2)\n", + " assert syntactic_eq(add(1, 3), add(2, 2))\n", "\n", "\n", "assert isinstance(add, Operation)\n", @@ -104,7 +104,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 3, "id": "6e293b33", "metadata": {}, "outputs": [], @@ -183,7 +183,7 @@ "\n", "\n", "def test_adding_a_variable():\n", - " assert add(x, 2) == add(x, 2)\n", + " assert syntactic_eq(add(x, 2), add(x, 2))\n", "\n", "\n", "fails(test_adding_a_variable)\n", @@ -374,7 +374,7 @@ "\n", "\n", "def test_mixed_together():\n", - " assert add(add(x, 2), add(2, y)) == add(x, add(add(3, y), 1))\n", + " assert syntactic_eq(add(add(x, 2), add(2, y)), add(x, add(add(3, y), 1)))\n", "\n", "\n", "fails(test_mixed_together)\n", @@ -512,7 +512,7 @@ "notebook_metadata_filter": "-all" }, "kernelspec": { - "display_name": "base", + "display_name": "effectful (3.12.9.final.0)", "language": "python", "name": "python3" }, @@ -526,7 +526,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.13" + "version": "3.12.9" } }, "nbformat": 4, diff --git a/effectful/handlers/jax/_handlers.py b/effectful/handlers/jax/_handlers.py index 4c021e9b6..ae70f4993 100644 --- a/effectful/handlers/jax/_handlers.py +++ b/effectful/handlers/jax/_handlers.py @@ -10,6 +10,7 @@ try: import jax + import jax.core import jax.numpy as jnp except ImportError: raise ImportError("JAX is required to use effectful.handlers.jax") @@ -33,6 +34,9 @@ IndexElement = None | int | slice | Sequence[int] | EllipsisType | jax.Array +class JaxOperation[**P, T](Operation[P, T]): ... + + def is_eager_array(x): return isinstance(x, jax.Array) or ( isinstance(x, Term) @@ -193,7 +197,7 @@ def _register_jax_op[**P, T](jax_fn: Callable[P, T]): if getattr(jax_fn, "__name__", None) == "__getitem__": return jax_getitem - @defop + @JaxOperation.define def _jax_op(*args, **kwargs) -> jax.Array: tm = defdata(_jax_op, *args, **kwargs) sized_fvs = sizesof(tm) @@ -210,6 +214,8 @@ def _jax_op(*args, **kwargs) -> jax.Array: ) ): raise NotHandled + elif _jax_op is jax_getitem and not args[1]: + return args[0] elif sized_fvs and set(sized_fvs.keys()) == fvsof(tm) - {jax_getitem, _jax_op}: # note: this cast is a lie. partial_eval can return non-arrays, as # can jax_fn. for example, some jax functions return tuples, @@ -281,7 +287,7 @@ def _einsum_named(subscripts, *operands, **kwargs) -> jax.Array: return jax.numpy.einsum(subscripts, *operands, **kwargs) # forward if any operand has a symbolic (Term) shape - if any(isinstance(arr.shape, Term) for arr in operands): + if any(isinstance(arr, Term) and not is_eager_array(arr) for arr in operands): raise NotHandled named = [_named_dims(op) for op in operands] @@ -367,6 +373,8 @@ def bind_dims[T, A, B]( >>> bind_dims(t, b, a).shape (3, 2) """ + if not names: + return value if isinstance(value, Term) and value.op == bind_dims: return bind_dims(value.args[0], *(names + tuple(value.args[1:]))) if jax.tree_util.treedef_is_leaf(jax.tree.structure(value)): @@ -374,6 +382,11 @@ def bind_dims[T, A, B]( return jax.tree.map(lambda v: bind_dims(v, *names), value) +@bind_dims.register(object) # type: ignore[attr-defined] +def _(*args, **kwargs): + raise NotHandled + + @defop @_CustomSingleDispatchCallable def unbind_dims[T, A, B]( @@ -382,11 +395,18 @@ def unbind_dims[T, A, B]( *names: Annotated[Operation[[], jax.Array], Scoped[B]], ) -> Annotated[T, Scoped[A | B]]: """Convert positional dimensions to named dimensions.""" + if not names: + return value if jax.tree_util.treedef_is_leaf(jax.tree.structure(value)): return __dispatch(typeof(value))(value, *names) return jax.tree.map(lambda v: unbind_dims(v, *names), value) +@unbind_dims.register(object) # type: ignore[attr-defined] +def _(*args, **kwargs): + raise NotHandled + + def jit(f, *args, **kwargs): f_noindex, f_reindex = _indexed_func_wrapper(f, jax_getitem, sizesof) f_noindex_jitted = jax.jit(f_noindex, *args, **kwargs) @@ -440,7 +460,18 @@ def _(x: jax.Array, other) -> bool: ) -@syntactic_hash.register(jax.Array) +@syntactic_eq.register +def _(x: jax.core.Tracer, other) -> bool: + return isinstance(other, jax.core.Tracer) and id(x) == id(other) + + +@syntactic_hash.register def _(x: jax.Array) -> int: # Concrete arrays aren't hashable; hash by shape, dtype, and bytes. - return hash(("jax.Array", x.shape, str(x.dtype), bytes(jax.numpy.asarray(x)))) + return hash(("jax.Array", x.shape, str(x.dtype), x.tobytes())) + + +@syntactic_hash.register +def _(x: jax.core.Tracer) -> int: + # Concrete arrays aren't hashable; hash by shape, dtype, and bytes. + return hash(("jax.core.Tracer", id(x))) diff --git a/effectful/handlers/jax/_terms.py b/effectful/handlers/jax/_terms.py index 53c4c094f..f29c36811 100644 --- a/effectful/handlers/jax/_terms.py +++ b/effectful/handlers/jax/_terms.py @@ -1,5 +1,6 @@ import functools import operator +import typing from collections.abc import Sequence from typing import Any, cast @@ -10,6 +11,7 @@ IndexElement, _register_jax_op, bind_dims, + is_eager_array, jax_getitem, unbind_dims, ) @@ -461,10 +463,10 @@ def _bind_dims_array(t: jax.Array, *args: Operation[[], jax.Array]) -> jax.Array return t # ensure that the result is a jax_getitem with an array as the first argument - if not (t.op is jax_getitem and isinstance(t.args[0], jax.Array)): + if not (t.op is jax_getitem and is_eager_array(t)): raise NotHandled - array = t.args[0] + array = typing.cast(jax.Array, t.args[0]) dims = t.args[1] assert isinstance(dims, Sequence) ndim = len(array.shape) diff --git a/effectful/handlers/jax/lax/__init__.py b/effectful/handlers/jax/lax/__init__.py new file mode 100644 index 000000000..8b2ea5c77 --- /dev/null +++ b/effectful/handlers/jax/lax/__init__.py @@ -0,0 +1,19 @@ +from typing import TYPE_CHECKING + +import jax.lax + +from effectful.handlers.jax._handlers import _register_jax_op + +for name, op in jax.lax.__dict__.items(): + wrapped_value = None + if callable(op): + wrapped_value = _register_jax_op(op) + else: + continue + + globals()[name] = wrapped_value + + +# Tell mypy about our wrapped functions. +if TYPE_CHECKING: + from jax.lax import * # noqa: F403 diff --git a/effectful/handlers/jax/monoid.py b/effectful/handlers/jax/monoid.py index 96b708e89..7cf369894 100644 --- a/effectful/handlers/jax/monoid.py +++ b/effectful/handlers/jax/monoid.py @@ -1,6 +1,9 @@ import functools import logging import typing +from collections import defaultdict +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass from typing import Protocol import jax @@ -8,64 +11,73 @@ import opt_einsum from opt_einsum import get_symbol +import effectful.handlers.jax.lax as lax import effectful.handlers.jax.numpy as jnp from effectful.handlers.jax import bind_dims, jax_getitem, unbind_dims from effectful.handlers.jax._handlers import is_eager_array from effectful.handlers.jax.scipy.special import logsumexp from effectful.ops.monoid import ( + And, CartesianProduct, + EvaluateIntp, Max, Min, Monoid, NormalizeIntp, + Or, Product, Streams, Sum, + _conjuncts, + _is_monoid_mask, _is_monoid_plus, - choose_contraction, + _is_simple_range, + complement, distributes_over, + is_equality, ) +from effectful.ops.monoid import Union as UnionM from effectful.ops.semantics import evaluate, fvsof, fwd, handler, typeof -from effectful.ops.syntax import ObjectInterpretation, deffn, implements -from effectful.ops.types import Interpretation, NotHandled, Operation, Term +from effectful.ops.syntax import ( + ObjectInterpretation, + as_dict, + deffn, + implements, + ite, +) +from effectful.ops.types import Expr, Interpretation, Operation, Term logger = logging.getLogger(__name__) -def cartesian_prod(x, y): - if x.ndim == 1: - x = x[:, None] - if y.ndim == 1: - y = y[:, None] - nx, dx = x.shape - ny, dy = y.shape - # Broadcast into (nx, ny, dx+dy), then flatten the first two axes - x_b = jnp.broadcast_to(x[:, None, :], (nx, ny, dx)) - y_b = jnp.broadcast_to(y[None, :, :], (nx, ny, dy)) - return jnp.concatenate([x_b, y_b], axis=-1).reshape(nx * ny, dx + dy) - - LogSumExp = Monoid(name="LogSumExp", identity=jnp.asarray(float("-inf"))) # ``Sum`` in log space is multiplication, which distributes over ``LogSumExp``: # a + logsumexp(b, c) = logsumexp(a + b, a + c) distributes_over.register(Sum, LogSumExp) +is_equality.register(jnp.equal) +for a, b in { + (jnp.less, jnp.greater), + (jnp.less_equal, jnp.greater_equal), + (jnp.equal, jnp.not_equal), +}: + assert isinstance(a, Operation) + assert isinstance(b, Operation) + complement.register(a, b) + def _jax_args(args): """True iff ``args`` is non-empty and every arg is a concrete - :class:`jax.typing.ArrayLike` or named tensor. At least one argument must be - a jax-related type. + :class:`jax.typing.ArrayLike` or named tensor. """ - return ( - bool(args) - and all(is_eager_array(a) or isinstance(a, jax.typing.ArrayLike) for a in args) - and any(is_eager_array(a) or isinstance(a, jax.Array) for a in args) + return args and all( + isinstance(a, jax.typing.ArrayLike) or is_eager_array(a) for a in args ) -class PlusJaxUpcast(ObjectInterpretation): +class PlusCastArray(ObjectInterpretation): @implements(Monoid.plus) def plus(self, monoid, *args): arg_types = [typeof(a) for a in args] @@ -127,32 +139,39 @@ def plus(self, *args): return functools.reduce(jnp.logaddexp, args) -class CartesianProductPlusJax(ObjectInterpretation): - @implements(CartesianProduct.plus) +class AndPlusJax(ObjectInterpretation): + @implements(And.plus) def plus(self, *args): - # Skip identity ``[()]`` args; short-circuit on zero ``[]``. Both - # sentinels arrive as Python lists alongside jax-array factors, so - # check for them explicitly before composing. - if not any(isinstance(a, jax.Array) for a in args): + if not _jax_args(args): return fwd() - result = None - for a in args: - if a is CartesianProduct.zero: - return CartesianProduct.zero - if a is CartesianProduct.identity: - continue - if not isinstance(a, jax.Array): - return fwd() - result = a if result is None else cartesian_prod(result, a) - if result is None: - return CartesianProduct.identity - # CartesianProduct values are streams of rows. ``cartesian_prod`` - # already lifts 1D inputs to 2D, but a single-array call seeds - # ``result = a`` unchanged — promote so the rank invariant holds for - # every array-path return. - if result.ndim == 1: - result = result[:, None] - return result + return functools.reduce(jnp.logical_and, args) + + +class OrPlusJax(ObjectInterpretation): + @implements(Or.plus) + def plus(self, *args): + if not _jax_args(args): + return fwd() + return functools.reduce(jnp.logical_or, args) + + +class IteJax(ObjectInterpretation): + @implements(ite) + def ite(self, cond, then, else_): + if not _jax_args((cond,)): + return fwd() + return jnp.where(cond, then, else_) + + +class MaskJax(ObjectInterpretation): + @implements(Monoid.mask) + def mask(self, monoid, value, mask): + if not ( + (is_eager_array(value) or not isinstance(value, Term)) + and (is_eager_array(mask) or not isinstance(mask, Term)) + ): + return fwd() + return jnp.where(mask, value, monoid.identity) class ReduceArrayGather(ObjectInterpretation): @@ -173,7 +192,7 @@ def reduce(self, monoid, body, streams): if typeof(body) is not jax.Array: return fwd() - if isinstance(body, Term) and body.op is delta: + if isinstance(body, Term) and body.op is monoid.delta: return fwd() body_fvs = fvsof(body) @@ -221,73 +240,151 @@ class ReduceArray(ObjectInterpretation): @implements(Monoid.reduce) def reduce(self, monoid, body, streams): - reductor = ARRAY_REDUCTORS.get(monoid, None) - if reductor is None: + if not is_eager_array(body): return fwd() - if typeof(body) is not jax.Array: + reductor = ARRAY_REDUCTORS.get(monoid, None) + if reductor is None: return fwd() - pos_dims = {} - if isinstance(body, Term): - if body.op == delta: - pos_dims = { - d.op - for d in body.args[0] - if isinstance(d, Term) and d.op in streams - } - elif _is_monoid_plus(body.op) and distributes_over( - body.op.__self__, monoid - ): - # delegate to factorization - return fwd() - body_fvs = fvsof(body) - used = { - k - for k, v in streams.items() - if k in body_fvs and k not in pos_dims and isinstance(v, range) - } + used = {k for k, v in streams.items() if k in body_fvs and isinstance(v, range)} if not used: return fwd() - delta_key = tuple(k() for k in streams if k in used) - arr = monoid.reduce(delta(delta_key, body), streams) + index = tuple(k() for k in streams if k in used) + arr = monoid.reduce(monoid.delta(index, body), streams) reduced_body = reductor(arr, axis=tuple(range(len(used)))) return reduced_body -@Operation.define -def delta(_index: tuple[int, ...], _weight: jax.Array) -> jax.Array: - raise NotHandled +class ReduceDisequalityMask(ObjectInterpretation): + """M.reduce(M.mask(v, And.plus(a != b, c)), S) + ≡ M.reduce(M.plus(M.mask(v, And.plus(a < b, c)), M.mask(v, And.plus(a > b, c)))) + """ + @implements(Monoid.reduce) + def _(self, monoid, body, streams: Streams): + match body: + case Term(mask_op, (value, mask), {}) if ( + _is_monoid_mask(mask_op) and mask_op.__self__ == monoid + ): + pass -def _range_stop(term: Term): - assert term.op == jnp.arange - if "stop" in term.kwargs: - return term.kwargs["stop"] - if len(term.args) < 2: - return term.args[0] - return term.args[1] + case _: + return fwd() + + mask_elems = _conjuncts(mask) + def _neq_to_plus(args, tail_mask_elems): + match args: + case (Term(stream_op, (), {}), index) if _is_simple_range( + streams.get(stream_op, None) + ): + pass + case _: + return None + return monoid.reduce( + monoid.plus( + monoid.mask(value, And.plus(stream_op() < index, *tail_mask_elems)), + monoid.mask(value, And.plus(stream_op() > index, *tail_mask_elems)), + ), + streams, + ) -class DeltaEmpty(ObjectInterpretation): - """delta((), weight) ≡ weight""" + for i, elem in enumerate(mask_elems): + match elem: + case Term(jnp.not_equal, args, {}): + tail_mask_elems = [e for (j, e) in enumerate(mask_elems) if i != j] + ret = _neq_to_plus(args, tail_mask_elems) or _neq_to_plus( + tuple(reversed(args)), tail_mask_elems + ) + if ret: + return ret - @implements(delta) - def _(self, index, weight): - if not index: - return weight return fwd() -class DeltaFusion(ObjectInterpretation): - """delta(i1, delta(i2, weight)) ≡ delta(i1 ++ i2, weight)""" +class ReduceArrayScan(ObjectInterpretation): + @implements(Monoid.reduce) + def _(self, monoid, body, streams: Streams): + match body: + case Term(mask_op, (value, mask), {}) if ( + _is_monoid_mask(mask_op) and mask_op.__self__ == monoid + ): + pass + + case _: + return fwd() - @implements(delta) - def _(self, index, weight): - if isinstance(weight, Term) and weight.op == delta: - return delta(index + weight.args[0], weight.args[1]) + def _inequality_to_scan(cmp_op, args, tail_mask_elems): + match args: + case (Term(stream_op, (), {}), index): + pass + case _: + return None + + stream = streams.get(stream_op, None) + if not _is_simple_range(stream): + return None + assert isinstance(stream, range) + + reverse = cmp_op in (jnp.greater_equal, jnp.greater) + n = len(stream) + pos_value = monoid.reduce( + monoid.delta((stream_op(),), value), {stream_op: stream} + ) + # inclusive prefix (or suffix, when reverse) scan over the stream + inclusive = lax.associative_scan(monoid.plus, pos_value, reverse=reverse) + # The strict comparisons (`<`, `>`) need an *exclusive* scan: pad + # an identity element on the appropriate side and shift the result + # so that scan_val[k] aggregates only the positions satisfying the + # comparison against k. The non-strict comparisons (`<=`, `>=`) use + # the inclusive scan directly. + match cmp_op: + case jnp.less: + scan_val = jnp.pad( + inclusive, + (1, 0), + mode="constant", + constant_values=monoid.identity, + )[:n] + case jnp.greater: + scan_val = jnp.pad( + inclusive, + (0, 1), + mode="constant", + constant_values=monoid.identity, + )[1:] + case _: + scan_val = inclusive + + tail_body = monoid.mask( + jax_getitem(scan_val, (index,)), And.plus(*tail_mask_elems) + ) + tail_streams = {k: v for (k, v) in streams.items() if k != stream_op} + if tail_streams: + return monoid.reduce(tail_body, tail_streams) + return tail_body + + mask_elems = _conjuncts(mask) + for i, elem in enumerate(mask_elems): + tail_mask_elems = [e for (j, e) in enumerate(mask_elems) if i != j] + match elem: + case Term( + ( + jnp.less_equal | jnp.greater_equal | jnp.less | jnp.greater + ) as cmp_op, + args, + {}, + ): + ret = _inequality_to_scan( + cmp_op, args, tail_mask_elems + ) or _inequality_to_scan( + complement.of(cmp_op), tuple(reversed(args)), tail_mask_elems + ) + if ret is not None: + return ret return fwd() @@ -302,7 +399,7 @@ class ReduceDeltaSimpleRange(ObjectInterpretation): @implements(Monoid.reduce) def _(self, monoid: Monoid, body, streams: Streams): - if not (isinstance(body, Term) and body.op == delta): + if not (isinstance(body, Term) and body.op == monoid.delta): return fwd() index, weight = body.args @@ -369,127 +466,17 @@ def _jax_getitem(arr, index): gathered_weight = handler(gather_subst)(evaluate)(sliced_weight) gathered_streams = handler(gather_subst)(evaluate)(sliced_streams) + gathered_body = ( + monoid.delta(tail_index, gathered_weight) if tail_index else gathered_weight + ) inner = ( - monoid.reduce(delta(tail_index, gathered_weight), gathered_streams) + monoid.reduce(gathered_body, gathered_streams) if gathered_streams else gathered_weight ) return bind_dims(inner, fresh_op) -class ReduceDependentRangeMask(ObjectInterpretation): - """Eliminate a dependent range by masking. - - reduce(M, streams ∪ {u: range(N), v: range(u())}, body) - ═══════════════════════════════════════════════════════════════════════════ - reduce(M, streams ∪ {u: range(N), v: range(N)}, where(v() < u(), body, M.identity)) - - Currently recognises only the lower-triangular form ``v: range(u())``: - constant start of 0, dependent stop equal to a bare call of another - stream var. - - Not yet supported: - - - **Upper-triangular** (``v: range(u(), N)`` — constant stop, dependent - start): bbox becomes ``range(0, N)`` (or ``range(0, bbox_N)``), guard - becomes ``v() >= u()``. Same shape of rewrite as lower-tri; differs - only in which side of the range carries the stream-var reference and - in the predicate direction. - - **Banded** (``v: range(u() - k, u() + k + 1)`` — two-sided dependent - bounds with constant width): bbox is ``range(0, N + k)`` (or similar - bounded by both endpoints' extents), guard is - ``(v() >= u() - k) & (v() < u() + k + 1)``. Needs both-sides - affine-bound recognition. - - **Strided dependent** (``v: range(0, u(), k)`` for ``k != 1``): bbox - stays ``range(0, N)`` and guard becomes - ``(v() < u()) & (v() % k == 0)`` (or equivalent), or alternatively - embed in a smaller bbox ``range(0, ceil(N/k))`` and remap the index. - - **Affine bounds** (``v: range(a*u() + b, c*u() + d)`` for affine - coefficients): bbox computed from ``ub(c*u() + d)`` over ``u``'s - range; guard is the conjunction of the two affine constraints. This - subsumes the upper/banded/strided cases under one affine recogniser. - - **Multi-stream-var dependent** (``v: range(u() + w())`` referencing - more than one outer stream var): bbox is the affine combination over - both referents' ranges; guard threads through all dependencies. - - **Reverse-order dependent ranges**: e.g. ``v: range(u(), 0, -1)``; - needs to handle negative step and the corresponding reverse - enumeration. - """ - - @implements(Monoid.reduce) - def _(self, monoid: Monoid, body, streams: Streams): - stream_vars = set(streams.keys()) - - # streams of the form k: range(X) - simple_ranges = { - k: v - for (k, v) in streams.items() - if isinstance(v, range) and v.start == 0 and v.step == 1 - } - for u, u_stream in simple_ranges.items(): - if fvsof(u_stream) & stream_vars: - continue - - for v, v_stream in streams.items(): - if ( - isinstance(v_stream, Term) - and v_stream.op == jnp.arange - and isinstance(_range_stop(v_stream), Term) - and _range_stop(v_stream).op == u - ): - fresh_streams = { - a: (u_stream if a == v else b) for (a, b) in streams.items() - } - - # there are other commuting rules for delta that we do not - # currently include - if isinstance(body, Term) and body.op == delta: - fresh_body = delta( - body.args[0], - jnp.where(v() < u(), body.args[1], monoid.identity), # type: ignore[arg-type] - ) - else: - fresh_body = jnp.where(v() < u(), body, monoid.identity) - - return monoid.reduce(fresh_body, fresh_streams) - - return fwd() - - -# Cross-cutting delta rules not yet implemented: -# -# - **Delta-commuting** (DC-hoist): for any pure op ``f`` (no Scoped binders -# that intersect a delta's index ops), push delta outward: -# f(args..., delta(idx, body), args...) -# ≡ delta(idx, f(args..., body, args...)) -# This normalizes delta to the outermost position so the reduce rules can -# pattern-match ``isinstance(body, Term) and body.op == delta`` cleanly. -# The soundness condition is mechanical via ``op.__fvs_rule__``: refuse to -# commute when a non-delta arg's scope binds any op in the delta's idx. -# -# - **Delta-merging** (DC-merge): under a pure binary op ``f`` (or -# generalized n-ary), merge multiple deltas when their index tuples are -# subsequence-compatible: -# f(delta(idx_a, v), delta(idx_b, w)) ≡ delta(idx_max, f(v, w)) -# where ``idx_max`` is the longer of ``idx_a``, ``idx_b`` and ``idx_a`` is -# a subsequence of ``idx_b`` (or vice versa). Refuse to fire when neither -# is a subsequence of the other, since that would silently insert an -# outer-product broadcast. - - -class ContractLongestArrayStream(ObjectInterpretation): - @implements(choose_contraction) - def _(self, factors, streams): - lengths = { - k: v.shape[0] if isinstance(v, jax.Array) and v.shape else 0 - for (k, v) in streams.items() - } - longest = max(lengths.values()) - return fwd( - factors, {k: v for (k, v) in streams.items() if lengths[k] == longest} - ) - - class ReduceSumProductContraction(ObjectInterpretation): """Fast-path a sum-of-products contraction.""" @@ -520,70 +507,306 @@ def _(self, body, streams: Streams): return fwd() # create leading reduction dimensions - delta_key = tuple(k() for k in streams) - pos_lhs = Sum.reduce(delta(delta_key, lhs), streams) - pos_rhs = Sum.reduce(delta(delta_key, rhs), streams) + index = tuple(k() for k in streams) + pos_lhs = Sum.reduce(Sum.delta(index, lhs), streams) + pos_rhs = Sum.reduce(Sum.delta(index, rhs), streams) dims = "".join(get_symbol(i) for i in range(len(streams))) contraction = jnp.einsum(f"{dims}...,{dims}...->...", pos_lhs, pos_rhs) return contraction -@jax.jit(static_argnums=(0,)) -def einsum(subscripts: str, /, *operands: jax.Array) -> jax.Array: - """Evaluate an einsum expression using monoid reductions.""" - if not operands: - raise ValueError("einsum requires at least one operand") +@dataclass +class Node: + ordinal: frozenset[str] + children: list["Node"] + factors: list[int] - in_spec, out_spec, _ = opt_einsum.parser.parse_einsum_input( - [subscripts, *(op.shape for op in operands)], shapes=True - ) - in_specs = in_spec.split(",") - all_letters = set(out_spec) | {c for s in in_specs for c in s} - ops = {c: Operation.define(jax.Array, name=c) for c in all_letters} +class _EinsumBuilder: + in_specs: Sequence[str] + out_spec: str + plates: frozenset[str] + operands: Sequence[jax.Array] + + def __init__( + self, subscripts: str, /, *operands: jax.Array, plates: str | None = None + ): + if not operands: + raise ValueError("einsum requires at least one operand") + + in_spec, out_spec, _ = opt_einsum.parser.parse_einsum_input( + [subscripts, *(op.shape for op in operands)], shapes=True + ) + in_specs = in_spec.split(",") + + self.in_specs = in_specs + self.out_spec = out_spec + self.plates = frozenset(plates or "") + self.operands = operands + + # check that the output spec preserves required plates + out_spec_set = frozenset(self.out_spec) + out_plates = out_spec_set & self.plates + for c in out_spec_set - self.plates: + missing_plates = self.ordinal[c] - out_plates + if missing_plates: + raise ValueError( + "It is nonsensical to preserve a plated dim without preserving " + f"all of that dim's plates, but found {c!r} without " + f"{','.join(sorted(missing_plates))!r}" + ) + + @functools.cached_property + def sizes(self) -> Mapping[str, int]: + """`sizes[d]` is the length of a dimension `d`.""" + sizes: dict[str, int] = {} + for spec, op in zip(self.in_specs, self.operands, strict=True): + for l, s in zip(spec, op.shape, strict=True): + if l in sizes and sizes[l] != s: + raise ValueError(f"Dimension {l} given sizes {s} and {sizes[l]}") + else: + sizes[l] = s + for c in self.out_spec: + if c not in sizes: + raise ValueError(f"einsum: output index {c!r} not present in any input") + return sizes + + @functools.cached_property + def ordinal(self) -> Mapping[str, frozenset[str]]: + """`ordinal[d]` is the plate context of a dimension `d`.""" + ordinal: dict[str, frozenset[str]] = defaultdict(lambda: self.plates) + for spec in self.in_specs: + spec_set = frozenset(spec) + for c in spec_set - self.plates: + ordinal[c] &= spec_set + return ordinal + + @functools.cached_property + def plate_tree(self) -> Node: + # 1. ordinal (plate context) of each factor + factor_ordinal = { + k: frozenset(spec) & self.plates for k, spec in enumerate(self.in_specs) + } + + # 2. node set: every factor ordinal, the global root ∅, and every pairwise + # intersection. The intersection of two ordinals is the deepest context + # that contains both, i.e. their common ancestor -- a shared plate MUST + # have a node to live at, or containment can't be a tree. Closing under + # ∩ materializes those frame nodes (finite lattice, so it terminates). + ordinals: set[frozenset[str]] = set(factor_ordinal.values()) | {frozenset()} + changed = True + while changed: + changed = False + for A in list(ordinals): + for B in list(ordinals): + if (A & B) not in ordinals: + ordinals.add(A & B) + changed = True + + nodes = {o: Node(ordinal=o, children=[], factors=[]) for o in ordinals} + + # 3. parent of o = the largest ordinal strictly inside o ("next context out"). + # Validity: that maximum must dominate EVERY other strict subset, i.e. the + # strict subsets form a chain. Two incomparable maximal subsets means o is a + # join over two separate branches -- not a tree -> raise. + for o in ordinals: + if not o: # root ∅ has no parent + continue + subs = [p for p in ordinals if p < o] # strict subsets, present as nodes + parent = max(subs, key=len) + offenders = [s for s in subs if not (s <= parent)] + if offenders: + raise ValueError( + f"factors do not form a plate tree: context {set(o)} sits above " + f"incomparable sub-plates {[set(s) for s in offenders]} (a join, not a nest)" + ) + nodes[parent].children.append(nodes[o]) + + # 4. hang each factor on its ordinal's node + for k, o in factor_ordinal.items(): + nodes[o].factors.append(k) + + return nodes[frozenset()] # the root + + def dim_type(self, dim: str) -> type: + return ( + int + if dim in self.plates + else Mapping[tuple, int] + if self.ordinal[dim] + else int + ) + + @functools.cached_property + def dim_op(self) -> Mapping[str, Operation]: + return { + dim: Operation.define(self.dim_type(dim), name=dim) + for dim in set("".join(self.in_specs)) + } + + @functools.cached_property + def out_vars(self) -> Mapping[str, Operation]: + return {c: Operation.define(jax.Array, name=f"out_{c}") for c in self.out_spec} + + @functools.cached_property + def global_enums(self) -> frozenset[str]: + return frozenset(c for c, o in self.ordinal.items() if not o) + + @functools.cached_property + def arrays(self) -> list[Term]: + return [Operation.define(jax.Array)() for _ in self.operands] + + def dim_index(self, dim: str) -> Expr: + out_globals = self.global_enums & set(self.out_spec) + + return ( + self.out_vars[dim]() + if dim in out_globals + else self.dim_op[dim]() + if dim in self.plates or not self.ordinal[dim] + else self.dim_op[dim]()[ + tuple(self.dim_op[p]() for p in sorted(self.ordinal[dim])) + ] + ) + + def out_mask(self, vars): + eqs = tuple(self.dim_index(p) == self.out_vars[p]() for p in vars) + return And.plus(*eqs) + + def _build_plate_reductions( + self, plate_tree: Node, parent_plates: frozenset[str] = frozenset() + ) -> Expr[jax.Array]: + + masked_factors = [] + for factor_idx in plate_tree.factors: + spec = self.in_specs[factor_idx] + factor = jax_getitem( + self.arrays[factor_idx], tuple(self.dim_index(d) for d in spec) + ) + # Only plated output dims are delta'd per factor. Global output dims are + # delta'd once by the outer ``out_mask`` in ``_einsum_expr``; masking + # them here too would check the same equality (e.g. ``out_b == b``) in + # every factor that mentions the dim as well as the outer mask. + factor_out_dims = frozenset( + d + for d in set(spec) & set(self.out_vars) - self.plates + if self.ordinal[d] + ) - sizes: dict[str, int] = {} - for spec, op in zip(in_specs, operands, strict=True): - for l, s in zip(spec, op.shape, strict=True): - if l in sizes and sizes[l] != s: - raise ValueError(f"Dimension {l} given sizes {s} and {sizes[l]}") + if factor_out_dims: + # Preserving a plated output dim requires preserving all its plates, + # as checked in ``__init__``. + factor_out_plates = plate_tree.ordinal & set(self.out_vars) + masked_factors.append( + ite( + self.out_mask(factor_out_plates), + Sum.mask(factor, self.out_mask(factor_out_dims)), + factor, + ) + ) else: - sizes[l] = s - for c in out_spec: - if c not in sizes: - raise ValueError(f"einsum: output index {c!r} not present in any input") - - arrays = [Operation.define(jax.Array) for _ in operands] - factors = [ - unbind_dims(arr(), *(ops[c] for c in spec)) - for arr, spec in zip(arrays, in_specs, strict=True) - ] - body = Product.plus(*factors) - - out_tuple = tuple(ops[c]() for c in out_spec) - streams = {op: range(sizes[c]) for c, op in ops.items()} + masked_factors.append(factor) + + child_reductions = ( + self._build_plate_reductions(n, parent_plates | plate_tree.ordinal) + for n in plate_tree.children + ) + product = Product.plus(*masked_factors, *child_reductions) + if plate_tree.ordinal: + plate_streams = { + self.dim_index(p).op: range(self.sizes[p]) + for p in sorted(plate_tree.ordinal - parent_plates) + } + return Product.reduce(product, plate_streams) + return product + + @functools.cached_property + def term(self) -> Expr[Callable]: + # one stream of per-plate-assignment rows for each plated sum dim + rows = {} + plated_enums = {c: o for c, o in self.ordinal.items() if o} + for c, o in plated_enums.items(): + ps = [(p, self.dim_op[p]) for p in sorted(o)] + delta_idx = tuple(op() for (_, op) in ps) + c_streams = {op: range(self.sizes[p]) for (p, op) in ps} + v = Operation.define(int, name=f"{c}_v") + rows[c] = CartesianProduct.reduce( + UnionM.reduce([as_dict((delta_idx, v()))], {v: range(self.sizes[c])}), + c_streams, + ) + + streams: Streams = { + self.dim_op[c]: range(self.sizes[c]) + for c in self.global_enums + if c not in self.out_spec + } | {self.dim_op[c]: r for c, r in rows.items()} + + dims = [(self.out_vars[c], self.sizes[c]) for c in self.out_spec] + + reductions = self._build_plate_reductions(self.plate_tree) + reduction = Sum.reduce(reductions, streams) if streams else reductions + return deffn( + bind_dims( + deffn(reduction, *(self.out_vars[c] for c in self.out_spec))( + *(unbind_dims(jnp.arange(d), v) for (v, d) in dims) + ), + *(v for (v, _) in dims), + ), + *(a.op for a in self.arrays), + ) + + +@jax.jit(static_argnums=(0,), static_argnames=("plates",)) +def einsum( + subscripts: str, /, *operands: jax.Array, plates: str | None = None +) -> jax.Array: + """Evaluate an einsum expression using monoid reductions. + + Generalizes :func:`jax.numpy.einsum` with plated dimensions in the style of + :func:`pyro.ops.contract.einsum`: indices in ``plates`` are plate + dimensions, and reductions along plates are product reductions. A sum + dimension that always appears together with a plate denotes a distinct + variable for each slice of that plate; its plate context (ordinal) is the + intersection of the plate sets of the inputs that mention it. When such a + dimension appears in the output it must be accompanied by all of its + plates. + + The expression is represented naively: each plated sum dimension ranges + over a :data:`CartesianProduct` stream of per-plate-assignment rows, and + each plated input is a :data:`Product` reduction over its plates. + """ + expr = _EinsumBuilder(subscripts, *operands, plates=plates).term + with handler(NormalizeIntp): - norm = deffn(Sum.reduce(delta(out_tuple, body), streams), *arrays) - result = norm(*operands) - assert isinstance(result, jax.Array) + norm_expr = evaluate(expr) + + assert CartesianProduct.reduce not in fvsof(norm_expr), ( + "failed to eliminate cartesian products" + ) + + with handler(EvaluateIntp), handler(NormalizeIntp): + assert callable(norm_expr) + result = norm_expr(*operands) + assert isinstance(result, jax.Array), "failed to fully evaluate" return result -NormalizeIntp.extend( - ReduceArray(), - ReduceSumProductContraction(), - ReduceArrayGather(), - ReduceDeltaSimpleRange(), - ReduceDependentRangeMask(), - DeltaEmpty(), - DeltaFusion(), +EvaluateIntp.extend( SumPlusJax(), ProductPlusJax(), MinPlusJax(), MaxPlusJax(), LogSumExpPlusJax(), - CartesianProductPlusJax(), - ContractLongestArrayStream(), - PlusJaxUpcast(), + AndPlusJax(), + OrPlusJax(), + IteJax(), + MaskJax(), + ReduceSumProductContraction(), + ReduceArray(), + ReduceDeltaSimpleRange(), + ReduceArrayScan(), + PlusCastArray(), ) + +NormalizeIntp.extend(ReduceArrayGather()) diff --git a/effectful/handlers/jax/numpy/__init__.py b/effectful/handlers/jax/numpy/__init__.py index 9556bdacf..34ec5a604 100644 --- a/effectful/handlers/jax/numpy/__init__.py +++ b/effectful/handlers/jax/numpy/__init__.py @@ -1,4 +1,5 @@ import types +import typing from typing import TYPE_CHECKING import jax.numpy @@ -46,6 +47,22 @@ einsum = Operation.define(_einsum_named) + +@Operation.define +def asarray(a, **kwargs) -> jax.Array: + import jax.core + + from effectful.ops.semantics import typeof + from effectful.ops.types import NotHandled, Term + + if isinstance(a, Term): + if issubclass(typeof(a), jax.Array | jax.core.Tracer) and not kwargs: + return typing.cast(jax.Array, a) + else: + raise NotHandled + return jax.numpy.asarray(a, **kwargs) + + # Tell mypy about our wrapped functions. if TYPE_CHECKING: from jax.numpy import * # type: ignore[assignment] # noqa: F403 diff --git a/effectful/ops/monoid.py b/effectful/ops/monoid.py index 067f29e55..11d58ff55 100644 --- a/effectful/ops/monoid.py +++ b/effectful/ops/monoid.py @@ -3,19 +3,26 @@ import itertools import operator import typing -from collections import Counter, UserDict, defaultdict -from collections.abc import Callable, Generator, Iterable, Mapping, Sequence +from collections import UserDict, defaultdict +from collections.abc import Callable, Generator, Iterable, Mapping, Sequence, Sized from dataclasses import dataclass from graphlib import TopologicalSorter from typing import Annotated, Any +import effectful.ops.syntax +from effectful.internals.runtime import interpreter from effectful.ops.semantics import coproduct, evaluate, fvsof, fwd, handler, typeof from effectful.ops.syntax import ( ObjectInterpretation, Scoped, + _MappingTerm, + _NumberTerm, + as_dict, defdata, deffn, implements, + ite, + range_, syntactic_eq, syntactic_hash, ) @@ -97,15 +104,15 @@ def inner_streams_first(streams: dict[Operation, Expr]) -> Iterable[Operation]: class Monoid[W]: """A monoid with ``plus`` and ``reduce`` :class:`Operation` s.""" - _name: str + __name__: str identity: W def __init__(self, identity: W, name: str): - self._name = name + self.__name__ = name self.identity = identity def __repr__(self): - return f"Monoid({self._name!r})" + return f"Monoid({self.__name__!r}, {self.identity!r})" def __eq__(self, other): return id(self) == id(other) @@ -153,6 +160,14 @@ def weighted[T]( """ raise NotHandled + @Operation.define + def mask(self, value: W, cond: Any = True) -> W: + raise NotHandled + + @Operation.define + def delta[K](self, index: K, weight: W) -> Mapping[K, W]: + raise NotHandled + class MonoidWithZero[T](Monoid[T]): zero: T @@ -168,10 +183,25 @@ def __init__(self, name: str, identity: T, zero: T): ArgMax = Monoid(name="ArgMax", identity=(Max.identity, None)) Sum = Monoid(name="Sum", identity=0) Product = MonoidWithZero(name="Product", identity=1, zero=0) -# CartesianProduct values are "two-level indexable" (rows × positions). The -# identity ``[()]`` is one row of zero positions (composing with it preserves -# shape); the zero ``[]`` is no rows (absorbs under product). -CartesianProduct = MonoidWithZero(name="CartesianProduct", identity=[()], zero=[]) +CartesianProduct: MonoidWithZero[Sequence[Mapping]] = MonoidWithZero( + name="CartesianProduct", identity=[{}], zero=[] +) +Union: Monoid[Sequence[Mapping]] = Monoid(name="Union", identity=[]) +And = MonoidWithZero(name="And", identity=True, zero=False) +Or = Monoid(name="Or", identity=False) + + +def _conjuncts(mask) -> Sequence[Term]: + """Return the conjuncts of an ``And`` mask as a flat tuple.""" + match mask: + case Term(And.plus, elems, {}): + return elems + case _: + return (mask,) + + +def _is_simple_range(obj) -> bool: + return isinstance(obj, range) and obj.start == 0 and obj.step == 1 @dataclass @@ -185,14 +215,19 @@ def __call__(self, t: T) -> bool: return t in self.elems -is_commutative = _ExtensiblePredicate({Max, Min, Sum, Product}) -is_idempotent = _ExtensiblePredicate({Max, Min}) +is_commutative = _ExtensiblePredicate({Max, Min, Sum, Product, And, Or}) +is_idempotent = _ExtensiblePredicate({Max, Min, And, Or}) @dataclass class _ExtensibleBinaryRelation[S, T]: tuples: set[tuple[S, T]] + def __init__(self, *args): + self.tuples = set() + for s, t in args: + self.register(s, t) + def register(self, s: S, t: T) -> None: self.tuples.add((s, t)) @@ -200,8 +235,31 @@ def __call__(self, s: S, t: T) -> bool: return (s, t) in self.tuples -distributes_over = _ExtensibleBinaryRelation( - {(Max, Min), (Min, Max), (Sum, Min), (Sum, Max), (Product, Sum)} +class _ExtensiblePartialInvolution[S](_ExtensibleBinaryRelation[S, S]): + def register(self, s: S, t: S) -> None: + for existing, image in self.tuples: + if existing == s and image != t: + raise ValueError(f"{s!r} already matched to {image!r}") + if existing == t and image != s: + raise ValueError(f"{t!r} already matched to {image!r}") + super().register(s, t) + super().register(t, s) + + def of(self, t: S) -> S | None: + for a, b in self.tuples: + if a == t: + return b + return None + + +distributes_over: _ExtensibleBinaryRelation[Monoid, Monoid] = _ExtensibleBinaryRelation( + (Max, Min), + (Min, Max), + (Sum, Min), + (Sum, Max), + (Product, Sum), + (CartesianProduct, Union), + (And, Or), ) @@ -223,6 +281,170 @@ def _is_monoid_weighted(op: Operation) -> bool: return isinstance(owner, Monoid) and op is owner.weighted +def _is_monoid_mask(op: Operation) -> bool: + """True if ``op`` is the ``mask`` operation of some :class:`Monoid`.""" + owner = getattr(op, "__self__", None) + return isinstance(owner, Monoid) and op is owner.mask + + +def _is_monoid_delta(op: Operation) -> bool: + """True if ``op`` is the ``delta`` operation of some :class:`Monoid`.""" + owner = getattr(op, "__self__", None) + return isinstance(owner, Monoid) and op is owner.delta + + +class GetitemDelta(ObjectInterpretation): + """M.delta(i, v)[q] ≡ M.mask(v, i == q)""" + + @implements(_MappingTerm.__getitem__) + def _(self, value, index): + if isinstance(value, Term) and _is_monoid_delta(value.op): + return value.op.__self__.mask(value.args[1], value.args[0] == index) + return fwd() + + +class MaskBool(ObjectInterpretation): + @implements(Monoid.mask) + def _(self, monoid, value, cond): + if isinstance(cond, bool): + return value if cond else monoid.identity + return fwd() + + +class MaskFusion(ObjectInterpretation): + """M.mask(M.mask(value, i1), i2) ≡ M.mask(value, And.plus(i1, i2))""" + + @implements(Monoid.mask) + def _(self, monoid, value, cond): + if ( + isinstance(value, Term) + and _is_monoid_mask(value.op) + and value.op.__self__ == monoid + ): + return monoid.mask(value.args[0], And.plus(value.args[1], cond)) + return fwd() + + +is_equality = _ExtensiblePredicate({_NumberTerm.__eq__}) + + +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)) + + 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. + + 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. + """ + + @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``.""" + + def test(op, stream_op, mask_key): + return ( + is_equality(op) + and stream_op in streams + and _is_simple_range(streams[stream_op]) + and not (fvsof(mask_key) & set(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 + + def _eliminate(self, monoid, value, mask, streams): + """Discharge one eliminable equality constraint via a gather, or return + ``None`` if no constraint is eliminable.""" + 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] + return monoid.reduce( + monoid.mask( + monoid.mask( + value, + And.plus(stream.start <= mask_key, mask_key < stream.stop), + ), + 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}, + ) + 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 + ) + + @implements(Monoid.reduce) + def _(self, monoid, body, streams): + if not isinstance(body, Term): + 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() + + +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``. + + A reduce-stream-independent condition gates the whole reduction uniformly, + so it can be lifted out: when ``c`` holds both sides are ``reduce(v, S)``, + and when it fails both are the identity (``reduce(identity, S) = identity``). + This holds for any monoid -- no commutativity or distributivity required. + """ + + @implements(Monoid.reduce) + def _(self, monoid, body, streams): + if not ( + streams + and isinstance(body, Term) + and _is_monoid_mask(body.op) + and body.op.__self__ == monoid + ): + return fwd() + value, cond = body.args + if fvsof(cond) & set(streams): + return fwd() + return monoid.mask(monoid.reduce(value, streams), cond) + + class PlusEmpty(ObjectInterpretation): """plus() = 0""" @@ -261,93 +483,82 @@ def is_nested_plus(x): class PlusDistr(ObjectInterpretation): - """x + (y * z) = x * y + x * z""" + """x * (y + z) = x * y + x * z""" @implements(Monoid.plus) def plus(self, monoid: Monoid, *args): - if any( - isinstance(x, Term) - and _is_monoid_plus(x.op) - and distributes_over(monoid, x.op.__self__) - for x in args - ): - non_terms = [] + if len(args) < 2: + return fwd() - # group terms by their monoid - by_monoid: dict[Monoid, list[Term]] = defaultdict(list) - for t in args: - if isinstance(t, Term) and _is_monoid_plus(t.op): - by_monoid[t.op.__self__].append(t) - else: - non_terms.append(t) - - # distribute over each group - progress = False - final_sum = [] - for m, terms in by_monoid.items(): - if ( - len(terms) > 1 - and distributes_over(monoid, m) - and not distributes_over(m, monoid) - ): - progress = True - term_args = (t.args for t in terms) - dist_terms = ( - monoid.plus(*args) for args in itertools.product(*term_args) + for i, a in enumerate(args): + if ( + isinstance(a, Term) + and _is_monoid_plus(a.op) + and distributes_over(monoid, (inner_monoid := a.op.__self__)) + and not distributes_over(inner_monoid, monoid) + ): + if i > 0: + return monoid.plus( + *args[: i - 1], + inner_monoid.plus( + *(monoid.plus(args[i - 1], x) for x in a.args) + ), + *args[i + 1 :], ) - final_sum.append(m.plus(*dist_terms)) else: - final_sum += terms - if progress: - return monoid.plus(*non_terms, *final_sum) + return monoid.plus( + inner_monoid.plus( + *(monoid.plus(x, args[i + 1]) for x in a.args) + ), + *args[i + 2 :], + ) return fwd() -class PlusConsecutiveDups(ObjectInterpretation): - """x ⊕ x ⊕ y = x ⊕ y""" +class PlusOrder(ObjectInterpretation): + """Normalize plus ordering for commutative monoids. + + x ⊕ y ⊕ x = x ⊕ x ⊕ y + + """ + + @staticmethod + def _term_sort_key(t: Term) -> tuple[int, int]: + return (syntactic_hash(t), id(t)) @implements(Monoid.plus) def plus(self, monoid, *args): - if not is_idempotent(monoid): + if not is_commutative(monoid): return fwd() - dedup_args = ( - args[i] - for i in range(len(args)) - if i == 0 or not syntactic_eq(args[i - 1], args[i]) + sorted_args = tuple( + sorted(range(len(args)), key=lambda i: self._term_sort_key(args[i])) ) - return fwd(monoid, *dedup_args) - + if sorted_args == tuple(range(len(args))): + return fwd() + return monoid.plus(*(args[i] for i in sorted_args)) -class PlusDups(ObjectInterpretation): - """x ⊕ y ⊕ x = x ⊕ y""" - @dataclass - class _HashableTerm: - term: Term +class PlusConsecutiveDups(ObjectInterpretation): + """Normalize duplicate arguments for idempotent monoids. - def __eq__(self, other): - return syntactic_eq(self, other) + x ⊕ x ⊕ y = x ⊕ y - def __hash__(self): - return syntactic_hash(self) + """ @implements(Monoid.plus) def plus(self, monoid, *args): - if not (is_idempotent(monoid) and is_commutative(monoid)): - return fwd() - - # elim dups - args_count = Counter(self._HashableTerm(t) for t in args) - if len(args_count) < len(args): - dedup_args = [] - for t in args: - ht = self._HashableTerm(t) - if ht in args_count: - dedup_args.append(t) - del args_count[ht] - return fwd(monoid, *dedup_args) - return fwd() + if not is_idempotent(monoid): + return fwd() + + dedup_args = tuple( + i + for i in range(len(args)) + if i == 0 or not syntactic_eq(args[i - 1], args[i]) + ) + if dedup_args == tuple(range(len(args))): + return fwd() + return monoid.plus(*(args[i] for i in dedup_args)) class ReducePartial(ObjectInterpretation): @@ -390,17 +601,14 @@ def reduce(self, monoid, body, streams): class ReduceSplit(ObjectInterpretation): - """Implements the identity - reduce(R, S, b1 + ... + bn) = reduce(R, S, b1) + ... + reduce(R, S, bn) - """ - @implements(Monoid.reduce) def reduce(self, monoid, body, streams): if not is_commutative(monoid): return fwd() - if isinstance(body, Term) and body.op is monoid.plus: - return monoid.plus(*(monoid.reduce(x, streams) for x in body.args)) - return fwd() + if not (isinstance(body, Term) and body.op is monoid.plus): + return fwd() + + return monoid.plus(*(monoid.reduce(a, streams) for a in body.args)) @Operation.define @@ -427,27 +635,48 @@ def choose_contraction(factors: Sequence[Any], streams: Streams) -> Operation: assert False, "expected at least one subset-minimal stream" -class ReduceFactorization(ObjectInterpretation): - """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 - v has no dependents among the remaining streams (so it can be innermost) - and F_rest is nonempty (universal variables stay in the outer core). - """ - +class Factor(ObjectInterpretation): @implements(Monoid.reduce) def reduce(self, monoid, body, 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 + v has no dependents among the remaining streams (so it can be innermost) + and F_rest is nonempty (universal variables stay in the outer core). + + A reduce-monoid mask wrapping the plus is handled too: + + reduce(M.mask(⊗(F_v ∪ F_rest), c), {v} ∪ S) + = reduce(M.mask(⊗F_rest ⊗ reduce(⊗F_v, {v}), c), S) + + Because the mask gates the whole product, every factor effectively depends + on the condition's variables; folding ``fvsof(c)`` into each factor's free + variables keeps those streams in the outer core (a stream the mask depends + on is treated as universal and never pulled into the inner reduce). The + mask therefore stays outside the inner reduce unchanged. Soundness relies + on ``M.identity`` annihilating the inner monoid's plus (the semiring zero), + so masking distributes over the inner product. + """ + if not (is_commutative(monoid) and isinstance(body, Term)): + return fwd() + + # Optionally peel an outer mask of the reduce monoid. + cond = None + plus_term = body + if _is_monoid_mask(body.op) and body.op.__self__ is monoid: + plus_term, cond = body.args + if not ( - is_commutative(monoid) - and isinstance(body, Term) - and _is_monoid_plus(body.op) - and distributes_over(body.op.__self__, monoid) + isinstance(plus_term, Term) + and _is_monoid_plus(plus_term.op) + and distributes_over(plus_term.op.__self__, monoid) ): return fwd() - inner = body.op.__self__ + inner = plus_term.op.__self__ stream_keys = set(streams) - factors = [(a, fvsof(a)) for a in body.args] + cond_fvs = fvsof(cond) if cond is not None else set() + factors = [(a, fvsof(a) | cond_fvs) for a in plus_term.args] # candidates: innermost-eligible (no remaining stream depends on v), # non-universal (some factor doesn't mention v) @@ -466,7 +695,7 @@ def reduce(self, monoid, body, streams): if len(eligible) == 1: inner_stream = next(iter(eligible)) else: - inner_stream = choose_contraction(body.args, eligible) + inner_stream = choose_contraction(plus_term.args, eligible) inner_factor_ids = frozenset( i for i, (_, fvs) in enumerate(factors) if inner_stream in fvs @@ -509,9 +738,53 @@ def reduce(self, monoid, body, 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 +class ReduceUnfactor(ObjectInterpretation): + """Undo one :class:`Factor` layer beneath a compatible product. + + This rule is deliberately kept out of ``NormalizeIntp``: together with + ``Factor`` it would form a rewrite cycle. It is used only while preparing + a candidate for cartesian-product inversion. + """ + + @implements(Monoid.reduce) + def reduce(self, monoid: Monoid, body, streams): + if not is_commutative(monoid): + return fwd() + if not (isinstance(body, Term) and _is_monoid_plus(body.op)): + return fwd() + + product_monoid = body.op.__self__ + if not distributes_over(product_monoid, monoid): + return fwd() + + for i, factor in enumerate(body.args): + if not (isinstance(factor, Term) and factor.op is monoid.reduce): + continue + + inner_body, inner_streams = factor.args + assert isinstance(inner_streams, Mapping) + + merged_streams = dict(streams) | dict(inner_streams) + + inner_factors = ( + inner_body.args + if isinstance(inner_body, Term) and inner_body.op is product_monoid.plus + else (inner_body,) + ) + return monoid.reduce( + product_monoid.plus( + *body.args[:i], *inner_factors, *body.args[i + 1 :] + ), + merged_streams, + ) + return fwd() + + class ReduceDistributeCartesianProduct(ObjectInterpretation): """Eliminates a reduce over a cartesian product. ∑_x₁ ∑_x₂ ... ∑_xₙ ∏_i f(xᵢ) = ∏_i ∑_xᵢ f(xᵢ) @@ -523,6 +796,11 @@ class ReduceDistributeCartesianProduct(ObjectInterpretation): = reduce(⨁, reduce(⨂, reduce(⨁, body2, {vv: body1}), S1), S2) where × is the cartesian product and ⨂ distributes over ⨁. + The body may also be a ``⨂``-plus of such reductions. Each reduction is + row-substituted independently, and the row positions determine which plate + variables are unified. Ordinary row-independent factors remain outside the + peeled plate reduction. + Note: This could be generalized to grouped inversion [2]. [1] Braz, Rd, Eyal Amir, and Dan Roth. "Lifted first-order @@ -532,74 +810,251 @@ class ReduceDistributeCartesianProduct(ObjectInterpretation): """ @implements(Monoid.reduce) - def reduce(self, sum_monoid: Monoid, sum_body, sum_streams): - if not (is_commutative(sum_monoid) and isinstance(sum_body, Term)): + def reduce(self, monoid: Monoid, body, streams): + if not isinstance(body, Term): return fwd() - # body is a product or multiplication of products - if _is_monoid_plus(sum_body.op) and distributes_over( - sum_body.op.__self__, sum_monoid + if not any( + isinstance(v, Term) and v.op == CartesianProduct.reduce + for v in streams.values() ): - prod_reduces = sum_body.args - else: - prod_reduces = [sum_body] - - products: list[tuple[Monoid, Callable, Operation, Term]] = [] - for prod_reduce in prod_reduces: - if not ( - isinstance(prod_reduce, Term) and _is_monoid_reduce(prod_reduce.op) - ): - return fwd() - prod_monoid: Monoid = prod_reduce.op.__self__ - prod_body = prod_reduce.args[0] - prod_streams = typing.cast(Mapping, prod_reduce.args[1]) - if not ( - distributes_over(prod_monoid, sum_monoid) - and (len(products) == 0 or products[-1][0] == prod_monoid) - ): - return fwd() + return fwd() - if len(prod_streams) > 1 or len(prod_streams) == 0: + # ``Factor`` may have moved product factors into nested additive + # reductions. Normalize the whole candidate so ReduceUnfactor can see + # and merge both stream bundles. This isolated interpretation cannot + # cycle with Factor. Build the candidate outside the active handler; + # otherwise invoking ``monoid.reduce`` here would recursively redispatch + # this rule. + with interpreter(CartesianProductNormalizeIntp): + candidate = monoid.reduce(body, streams) + body, streams = candidate.args + + inner_reduces: tuple[Term, ...] + if isinstance(body, Term) and _is_monoid_reduce(body.op): + inner_reduces = (body,) + outer_factors: tuple = () + inner_monoid = body.op.__self__ + elif isinstance(body, Term) and _is_monoid_plus(body.op): + inner_monoid = body.op.__self__ + inner_reduces = tuple( + arg + for arg in body.args + if isinstance(arg, Term) and arg.op is inner_monoid.reduce + ) + if not inner_reduces: return fwd() - (prod_op, prod_stream) = next(iter(prod_streams.items())) - products.append( - (prod_monoid, deffn(prod_body, prod_op), prod_op, prod_stream) + outer_factors = tuple( + arg + for arg in body.args + if not (isinstance(arg, Term) and arg.op is inner_monoid.reduce) ) + else: + return fwd() + + if not distributes_over(inner_monoid, monoid): + return fwd() + + class InvalidIndexError(Exception): ... - assert len(products) > 0 + def drop_elem(ls, index): + return tuple(x for (i, x) in enumerate(ls) if i != index) - for outer_sum_streams, cprod_op, cprod_term in inner_stream(sum_streams): + for stream_key, stream_body in streams.items(): + # stream is cartesian if not ( - isinstance(cprod_term, Term) - and cprod_term.op is CartesianProduct.reduce + isinstance(stream_body, Term) + and stream_body.op is CartesianProduct.reduce ): continue - (cprod_body, cprod_streams) = cprod_term.args + # Product reductions that do not use this row are ordinary outer + # factors. Every factor that does use it must be plate-reduced. + row_inner_reduces = tuple( + reduce for reduce in inner_reduces if stream_key in fvsof(reduce) + ) + row_outer_factors = outer_factors + tuple( + reduce for reduce in inner_reduces if stream_key not in fvsof(reduce) + ) + if not row_inner_reduces or stream_key in fvsof(row_outer_factors): + continue + + (cprod_body, cprod_streams) = stream_body.args + assert isinstance(cprod_streams, dict) + + # plates are rectangular if not all( - prod_stream.op == cprod_op for (_, _, _, prod_stream) in products + isinstance(plate_stream, range) + for plate_stream in cprod_streams.values() ): continue - prod_op = Operation.define(products[0][2]) - prod_monoid = products[0][0] - inner_sum = sum_monoid.reduce( - prod_monoid.plus( - *(prod_body(prod_op()) for (_, prod_body, _, _) in products) - ), - {prod_op: cprod_body}, + # stream body is a sequence of mappings from plate index to domain value + match cprod_body: + case Term( + Union.reduce, + ( + [Term(effectful.ops.syntax.as_dict, ((idx, union_body),), {})], + union_streams, + ), + {}, + ) if isinstance(idx, Sequence) and set( + i.op for i in idx if isinstance(i, Term) + ) >= set(cprod_streams): + pass + case _: + continue + + assert len(idx) > 0 + + # inner product folds over all plates + plate_index, plate_op = next( + (j, i.op) for (j, i) in enumerate(idx) if i.op in cprod_streams ) - prod = prod_monoid.reduce(inner_sum, cprod_streams) - outer_sum = ( - sum_monoid.reduce(prod, outer_sum_streams) - if outer_sum_streams - else prod + plate_range = cprod_streams[plate_op] + + def row_substitute(inner_body, inner_streams): + """Peel ``plate_index`` off every ``stream_key[...]`` in one + summand. Asserts the summand's bundle contains a stream that + folds over the full cartesian product row and returns that plate + variable alongside the substituted body.""" + if stream_key in fvsof(inner_streams): + raise InvalidIndexError() + + inner_plate_op = None + + # substitute all instances of row[i, *rest] -> row[*rest] + def _getitem(mapping, idx1): + nonlocal inner_plate_op + + idx1 = idx1 if isinstance(idx1, Sequence) else (idx1,) + if isinstance(mapping, Term) and mapping.op == stream_key: + if not ( + isinstance(idx1, Sequence) + and len(idx1) > plate_index + and isinstance(idx1[plate_index], Term) + and idx1[plate_index].op in inner_streams + and syntactic_eq( + inner_streams[idx1[plate_index].op], plate_range + ) + ): + raise InvalidIndexError() + + if inner_plate_op is None: + inner_plate_op = idx1[plate_index].op + elif inner_plate_op != idx1[plate_index].op: + raise InvalidIndexError() + + return fwd(mapping, drop_elem(idx1, plate_index)) + return fwd() + + subst = handler({_MappingTerm.__getitem__: _getitem})(evaluate)( + inner_body + ) + if inner_plate_op is None and inner_streams: + # A nontrivial reduction that does not fold over the row + # cannot be unified with the peeled plate. Streamless + # reductions are ordinary product factors introduced by + # sum-of-products expansion and are retained unchanged. + raise InvalidIndexError() + return subst, inner_plate_op + + try: + substituted = [ + row_substitute(*reduce.args) for reduce in row_inner_reduces + ] + except InvalidIndexError: + continue + + # Unify reductions by the variable used at this row position, not + # merely by equal ranges. Equal-sized axes are otherwise ambiguous. + shared_plate_op = plate_op + combined_factors = [] + for (subst_body, inner_plate_op), inner_reduce in zip( + substituted, row_inner_reduces + ): + _, inner_streams = inner_reduce.args + assert isinstance(inner_streams, Mapping) + assert inner_plate_op is not None + if inner_plate_op is not shared_plate_op: + subst_body = handler({inner_plate_op: shared_plate_op})(evaluate)( + subst_body + ) + + inner_tail_streams = { + k: v for (k, v) in inner_streams.items() if k != inner_plate_op + } + if inner_tail_streams: + subst_body = inner_monoid.reduce(subst_body, inner_tail_streams) + combined_factors.append(subst_body) + + combined = ( + combined_factors[0] + if len(combined_factors) == 1 + else inner_monoid.plus(*combined_factors) ) - return outer_sum + + peeled_idx = drop_elem(idx, plate_index) + peeled_cprod_streams = { + k: v for (k, v) in cprod_streams.items() if k != plate_op + } + if not peeled_cprod_streams and not peeled_idx: + + def _to_body(mapping, key): + if isinstance(mapping, Term) and mapping.op == stream_key: + return union_body + return fwd() + + subst_combined = handler({_MappingTerm.__getitem__: _to_body})( + evaluate + )(combined) + inner_reduce_body = monoid.reduce(subst_combined, union_streams) + else: + peeled_body = Union.reduce( + [as_dict((peeled_idx, union_body))], union_streams + ) + if not peeled_cprod_streams: + peeled_cprod = peeled_body + else: + peeled_cprod = CartesianProduct.reduce( + peeled_body, peeled_cprod_streams + ) + inner_reduce_body = monoid.reduce(combined, {stream_key: peeled_cprod}) + + peeled_reduce = inner_monoid.reduce( + inner_reduce_body, + {shared_plate_op: plate_range}, + ) + + result_body = ( + inner_monoid.plus(peeled_reduce, *row_outer_factors) + if row_outer_factors + else peeled_reduce + ) + + # Include any extra sum streams outermost. In particular, the + # non-reduce product factors above remain outside the plate fold. + tail_streams = {k: v for (k, v) in streams.items() if k != stream_key} + if tail_streams: + result = monoid.reduce(result_body, tail_streams) + else: + result = result_body + + return result return fwd() +class ReduceUnion(ObjectInterpretation): + @implements(Monoid.reduce) + def reduce(self, monoid, body, streams): + for k, v in streams.items(): + if isinstance(v, Term) and v.op == Union.reduce: + union_body, union_streams = v.args + return monoid.reduce(body, streams | {k: union_body} | union_streams) + return fwd() + + class ReduceWeightedStream(ObjectInterpretation): """reduce(M, body, {x: WM.weighted(s, v, w), ...}) = reduce(M, WM.plus(w[v:=x()], body), {x: s, ...}) @@ -624,54 +1079,6 @@ def reduce(self, monoid, body, streams): return fwd() -class ReduceCartesianWeightedStream(ObjectInterpretation): - """``CartesianProduct.reduce`` over a :func:`weighted` body whose - ``weight`` is independent of the plate (product-index) streams:: - - CartesianProduct.reduce(M.weighted(s, w), plates) - = M.weighted( - CartesianProduct.reduce(s, plates), - deffn(M.reduce(w, {e: row()}), row), - ) - - Reuses ``body``'s element binder ``e`` (already typed by construction); - introduces a fresh ``row`` binder typed as ``Iterable[elem_type]``. - - Only fires when ``w`` is independent of the plate vars. - """ - - @Operation.define - @staticmethod - def _iterable_elem[T](iter: Iterable[T]) -> T: - raise NotHandled - - @implements(Monoid.reduce) - def reduce(self, monoid, body, streams): - if monoid is not CartesianProduct: - return fwd() - if not (isinstance(body, Term) and _is_monoid_weighted(body.op)): - return fwd() - - s, w = body.args - if not isinstance(s, Term) and len(s) == 0: - return CartesianProduct.reduce([], streams) - - if set(streams.keys()) & fvsof(w): - return fwd() - - elem_typ = typeof(self._iterable_elem(s)) - elem_op = Operation.define(elem_typ, name="elem") - row_op = Operation.define(Iterable[elem_typ], name="row") - - weight_monoid = body.op.__self__ - joint_weight = deffn( - weight_monoid.reduce(w(elem_op()), {elem_op: row_op()}), row_op - ) - joint_stream = CartesianProduct.reduce(s, streams) - - return weight_monoid.weighted(joint_stream, joint_weight) - - class MonoidOverCallable(ObjectInterpretation): """``monoid.reduce(f, streams) = lambda *a: monoid.reduce(f(*a), streams)``.""" @@ -734,6 +1141,14 @@ def _scalar_args(args): ) +class DeltaConcrete(ObjectInterpretation): + @implements(Monoid.delta) + def _(self, _, k, v): + if not fvsof((k, v)): + return {k: v} + return fwd() + + class SumPlus(ObjectInterpretation): """Scalar implementation of :data:`Sum`.""" @@ -802,6 +1217,18 @@ def plus(self, *args): return max(args, key=lambda a: a[0]) +def _disjoint_merge[K, V](*dicts: Mapping[K, V]) -> Mapping[K, V]: + merged = {} + for d in dicts: + if isinstance(d, Term): + return fwd() + for key, value in d.items(): + if key in merged: + raise ValueError(f"Duplicate key found: '{key}'") + merged[key] = value + return merged + + class CartesianProductPlus(ObjectInterpretation): """Pure-Python implementation of :data:`CartesianProduct`.""" @@ -813,16 +1240,22 @@ def plus(self, *args): return fwd() if not all(isinstance(x, Iterable) for x in args): return fwd() + return [_disjoint_merge(*vals) for vals in itertools.product(*args)] - def to_tuple(x): - return x if isinstance(x, tuple) else (x,) - return [ - sum((to_tuple(v) for v in 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)) -is_scalar = _ExtensiblePredicate({Min, Max, Sum, Product}) +is_scalar = _ExtensiblePredicate({Min, Max, Sum, Product, And, Or}) class MonoidOverSequence(ObjectInterpretation): @@ -906,7 +1339,6 @@ def reduce(self, monoid, body, streams): if not isinstance(vs, Term) and isinstance(vs, collections.abc.Sequence) and len(vs) == 1 - and isinstance(vs[0], Term) } if not singletons: return fwd() @@ -923,6 +1355,356 @@ def reduce(self, monoid, body, streams): return monoid.reduce(new_body, new_streams) if new_streams else new_body +class WhereHoist(ObjectInterpretation): + """Hoist :func:`ite` out of monoid ``reduce`` and ``plus``. + + A stream-independent selection commutes with reduction, while monoid + addition distributes pointwise over either selected branch. + """ + + @implements(Monoid.plus) + def plus(self, monoid, *args): + if len(args) < 2: + return fwd() + + for i, arg in enumerate(args): + if not ( + isinstance(arg, Term) + and arg.op is ite + and len(arg.args) == 3 + and not arg.kwargs + ): + continue + + cond, when_true, when_false = arg.args + return ite( + cond, + monoid.plus(*args[:i], when_true, *args[i + 1 :]), + monoid.plus(*args[:i], when_false, *args[i + 1 :]), + ) + + return fwd() + + @implements(Monoid.reduce) + def reduce(self, monoid, body, streams): + if not ( + streams + and isinstance(body, Term) + and body.op is ite + and len(body.args) == 3 + and not body.kwargs + ): + return fwd() + + cond, when_true, when_false = body.args + if fvsof(cond) & set(streams): + return fwd() + + return ite( + cond, + monoid.reduce(when_true, streams), + monoid.reduce(when_false, streams), + ) + + +class ReduceWhereEqualityPeel(ObjectInterpretation): + """Peel stream-independent conjuncts off a ``where`` equality guard. + + Given a condition ``outer & inner`` where ``outer`` is independent of the + reduced streams and ``inner`` contains an equality selecting one of them:: + + M.reduce(ite(outer & inner, x, y), S) + == ite(outer, + M.reduce(ite(inner, x, y), S), + M.reduce(y, S)) + + This exposes the stream equality to gather/elimination rules while moving + the remaining output guard above the reduction. + """ + + @staticmethod + def _stream_equality(cond, streams): + def matches(op, stream_term, other): + return ( + is_equality(op) + and isinstance(stream_term, Term) + and not stream_term.args + and not stream_term.kwargs + and stream_term.op in streams + and not (fvsof(other) & set(streams)) + ) + + return ( + isinstance(cond, Term) + and len(cond.args) == 2 + and ( + matches(cond.op, cond.args[0], cond.args[1]) + or matches(cond.op, cond.args[1], cond.args[0]) + ) + ) + + @implements(Monoid.reduce) + def reduce(self, monoid, body, streams): + if not ( + streams + and isinstance(body, Term) + and body.op is ite + and len(body.args) == 3 + and not body.kwargs + ): + return fwd() + + cond, when_true, when_false = body.args + if not (isinstance(cond, Term) and cond.op is And.plus): + return fwd() + + stream_ops = set(streams) + dependent = [c for c in cond.args if fvsof(c) & stream_ops] + independent = [c for c in cond.args if not (fvsof(c) & stream_ops)] + if not independent or not any( + self._stream_equality(c, streams) for c in dependent + ): + return fwd() + + return ite( + And.plus(*independent), + monoid.reduce(ite(And.plus(*dependent), when_true, when_false), streams), + monoid.reduce(when_false, streams), + ) + + +complement: _ExtensiblePartialInvolution[Operation] = _ExtensiblePartialInvolution( + (_NumberTerm.__ne__, _NumberTerm.__eq__), +) + + +class ReduceWhereToMasks(ObjectInterpretation): + """Split an equality-guarded ``where`` reduction into masked reductions. + + For a conjunction of stream-dependent equalities ``eqs``:: + + M.reduce(where(eqs, a, b), S) + == M.plus(M.reduce(M.mask(a, eqs), S), + M.reduce(M.mask(b, not(eqs)), S)) + + De Morgan's law represents ``not(eqs)`` as the disjunction of the + corresponding disequalities. The two masks are complementary and hence + partition the stream assignments without double counting. + """ + + @staticmethod + def _combine(op, terms): + return terms[0] if len(terms) == 1 else op(*terms) + + @implements(Monoid.reduce) + def reduce(self, monoid, body, streams): + if not (isinstance(body, Term) and body.op == ite): + return fwd() + + cond, when_true, when_false = body.args + conds = _conjuncts(cond) + all_have_compl = all( + isinstance(t, Term) and complement.of(t.op) is not None for t in conds + ) + stream_ops = set(streams) + all_stream_dep = all(bool(fvsof(t) & stream_ops) for t in conds) + if not (conds and all_have_compl and all_stream_dep): + return fwd() + + disequalities = tuple(complement.of(t.op)(*t.args, **t.kwargs) for t in conds) + return monoid.plus( + monoid.reduce(monoid.mask(when_true, cond), streams), + monoid.reduce( + monoid.mask(when_false, self._combine(Or.plus, disequalities)), + streams, + ), + ) + + +class ReduceDisjunctiveDisequalityMask(ObjectInterpretation): + """Partition a disjunctive disequality mask into disjoint reductions. + + For example, the overlapping regions ``i != x`` and ``j != y`` are + rewritten as ``i != x`` and ``i == x and j != y``:: + + reduce(mask(v, (i != x) or (j != y)), streams) + == plus( + reduce(mask(v, i != x), streams), + reduce(mask(v, (i == x) and (j != y)), streams), + ) + + More generally, disjunct ``k`` is conjoined with the complements of all + preceding disjuncts. The resulting masks are pairwise disjoint, so their + reductions can be combined with the reduction monoid. In particular, + each result has a conjunctive mask that :class:`ReduceArrayScan` can + eliminate. + """ + + @implements(Monoid.reduce) + def reduce(self, monoid, body, streams: Streams): + match body: + case Term(mask_op, (value, Term(Or.plus, disjuncts, {})), {}) if ( + _is_monoid_mask(mask_op) and mask_op.__self__ == monoid + ): + pass + case _: + return fwd() + + if len(disjuncts) < 2 or not all( + isinstance(disjunct, Term) and is_equality(complement.of(disjunct.op)) + for disjunct in disjuncts + ): + return fwd() + + preceding_complements: list = [] + reductions = [] + for disjunct in disjuncts: + assert isinstance(disjunct, Term) + reductions.append( + monoid.reduce( + monoid.mask(value, And.plus(*preceding_complements, disjunct)), + streams, + ) + ) + comp = complement.of(disjunct.op) + assert comp is not None + preceding_complements.append(comp(*disjunct.args, **disjunct.kwargs)) + + return monoid.plus(*reductions) + + +class ReduceDependentRangeMask(ObjectInterpretation): + """Eliminate a dependent range by masking. + + reduce(M, streams ∪ {u: range(N), v: range(u())}, body) + ═══════════════════════════════════════════════════════════════════════════ + reduce(M, streams ∪ {u: range(N), v: range(N)}, where(v() < u(), body, M.identity)) + + Currently recognises only the lower-triangular form ``v: range(u())``: + constant start of 0, dependent stop equal to a bare call of another + stream var. + + Not yet supported: + + - **Upper-triangular** (``v: range(u(), N)`` — constant stop, dependent + start): bbox becomes ``range(0, N)`` (or ``range(0, bbox_N)``), guard + becomes ``v() >= u()``. Same shape of rewrite as lower-tri; differs + only in which side of the range carries the stream-var reference and + in the predicate direction. + - **Banded** (``v: range(u() - k, u() + k + 1)`` — two-sided dependent + bounds with constant width): bbox is ``range(0, N + k)`` (or similar + bounded by both endpoints' extents), guard is + ``(v() >= u() - k) & (v() < u() + k + 1)``. Needs both-sides + affine-bound recognition. + - **Strided dependent** (``v: range(0, u(), k)`` for ``k != 1``): bbox + stays ``range(0, N)`` and guard becomes + ``(v() < u()) & (v() % k == 0)`` (or equivalent), or alternatively + embed in a smaller bbox ``range(0, ceil(N/k))`` and remap the index. + - **Affine bounds** (``v: range(a*u() + b, c*u() + d)`` for affine + coefficients): bbox computed from ``ub(c*u() + d)`` over ``u``'s + range; guard is the conjunction of the two affine constraints. This + subsumes the upper/banded/strided cases under one affine recogniser. + - **Multi-stream-var dependent** (``v: range(u() + w())`` referencing + more than one outer stream var): bbox is the affine combination over + both referents' ranges; guard threads through all dependencies. + - **Reverse-order dependent ranges**: e.g. ``v: range(u(), 0, -1)``; + needs to handle negative step and the corresponding reverse + enumeration. + """ + + @implements(Monoid.reduce) + def _(self, monoid: Monoid, body, streams: Streams): + for u, u_stream in streams.items(): + # streams of the form k: range(X) + if not _is_simple_range(u_stream): + continue + + for v, v_stream in streams.items(): + if not (isinstance(v_stream, Term) and v_stream.op == range_): + continue + + start, stop, step = v_stream.start, v_stream.stop, v_stream.step # type: ignore[attr-defined] + if not ( + isinstance(start, int) + and start == 0 + and isinstance(step, int) + and step == 1 + and isinstance(stop, Term) + and stop.op == u + ): + continue + + fresh_streams = { + a: (u_stream if a == v else b) for (a, b) in streams.items() + } + fresh_body = monoid.mask(body, v() < u()) + return monoid.reduce(fresh_body, fresh_streams) + + return fwd() + + +class ReduceDisequalityMask(ObjectInterpretation): + """M.reduce(M.mask(v, And.plus(a != b, c)), S) + ≡ M.reduce(M.plus(M.mask(v, And.plus(a < b, c)), M.mask(v, And.plus(a > b, c)))) + """ + + @implements(Monoid.reduce) + def _(self, monoid, body, streams: Streams): + match body: + case Term(mask_op, (value, mask), {}) if ( + _is_monoid_mask(mask_op) and mask_op.__self__ == monoid + ): + pass + + case _: + return fwd() + + mask_elems = _conjuncts(mask) + + def _neq_to_plus(args, tail_mask_elems): + match args: + case (Term(stream_op, (), {}), index) if _is_simple_range( + streams.get(stream_op, None) + ): + pass + case _: + return None + return monoid.reduce( + monoid.plus( + monoid.mask(value, And.plus(stream_op() < index, *tail_mask_elems)), + monoid.mask(value, And.plus(stream_op() > index, *tail_mask_elems)), + ), + streams, + ) + + for i, elem in enumerate(mask_elems): + if not (isinstance(elem, Term) and is_equality(complement.of(elem.op))): + continue + + tail_mask_elems = [e for (j, e) in enumerate(mask_elems) if i != j] + ret = _neq_to_plus(elem.args, tail_mask_elems) or _neq_to_plus( + tuple(reversed(elem.args)), tail_mask_elems + ) + if ret is not None: + return ret + + return fwd() + + +class ContractLongestStream(ObjectInterpretation): + @implements(choose_contraction) + def _(self, factors, streams): + lengths = { + k: len(v) if not isinstance(v, Term) and isinstance(v, Sized) else 0 + for (k, v) in streams.items() + } + longest = max(lengths.values()) + longest_streams = {k: v for (k, v) in streams.items() if lengths[k] == longest} + if len(longest_streams) == len(streams): + return fwd() + return choose_contraction(factors, longest_streams) + + class _ExtensibleInterpretation(UserDict, Interpretation): def extend(self, *intps: Interpretation) -> typing.Self: for intp in intps: @@ -930,34 +1712,183 @@ def extend(self, *intps: Interpretation) -> typing.Self: return self -NormalizeIntp = _ExtensibleInterpretation().extend( +EvaluateIntp = _ExtensibleInterpretation().extend( ReducePartial(), - EliminateSingletonStreams(), + DeltaConcrete(), + SumPlus(), + MinPlus(), + MaxPlus(), + ProductPlus(), + ArgMinPlus(), + ArgMaxPlus(), + CartesianProductPlus(), + UnionPlus(), + ReduceEqualityMaskRange(), + ReduceWhereToMasks(), +) + +CartesianProductNormalizeIntp = functools.reduce( + coproduct, + typing.cast( + tuple[Interpretation, ...], + ( + PlusEmpty(), + PlusSingle(), + PlusAssoc(), + ReduceUnfactor(), + ), + ), +) +"""Structural preprocessing used exclusively for cartesian-product inversion. + +It intentionally excludes ``Factor`` and sum-of-products rewrites, which +would either cycle with ``ReduceUnfactor`` or move factors across a product +fold. +""" + + +NormalizeIntp = _ExtensibleInterpretation().extend( + GetitemDelta(), MonoidOverSequence(), MonoidOverMapping(), MonoidOverCallable(), ReduceFusion(), + ReduceUnion(), ReduceSplit(), - ReduceFactorization(), + Factor(), ReduceDistributeCartesianProduct(), ReduceWeightedStream(), - ReduceCartesianWeightedStream(), + ReduceMaskHoist(), + EliminateSingletonStreams(), PlusEmpty(), PlusSingle(), PlusAssoc(), PlusDistr(), PlusConsecutiveDups(), - PlusDups(), - SumPlus(), - MinPlus(), - MaxPlus(), - ProductPlus(), - ArgMinPlus(), - ArgMaxPlus(), - CartesianProductPlus(), + PlusOrder(), PlusCastFloat(), + MaskFusion(), + MaskBool(), + WhereHoist(), + ReduceWhereEqualityPeel(), + ReduceDisjunctiveDisequalityMask(), + ReduceDependentRangeMask(), + ReduceDisequalityMask(), + ContractLongestStream(), ) -"""``NormalizeIntp``applies pure-Term rewrites (associativity, distributivity, -identity elimination, fusion, factorization, etc.). +"""``NormalizeIntp`` applies pure-Term rewrites (associativity, distributivity, +identity elimination, fusion, factorization, etc.) that drive a reduce +expression toward a *normal form*. + +Normal form +=========== + +The rules collectively push expressions toward an einsum / variable-elimination +shape: a factored sum-of-products over independent index ranges, with weights +carried symbolically by ``mask`` and ``delta`` rather than lowered to concrete +array ops. Throughout, "sum" means a ``reduce``/``plus`` of an additive monoid +``R`` and "product" a ``plus`` of a multiplicative monoid ``M`` with +``distributes_over(M, R)``. + +A fully normalized leaf expression matches the grammar (for additive ``R`` / +multiplicative ``M``):: + + nf ::= container[nf] # lambdas/dicts/sequences outermost + | M.mask(prod, orphan_cond) # only conjuncts no factor mentions + | prod + prod ::= M.plus(factor, ...) # flattened, n-ary + | factor + factor ::= M.mask(core, cond) # conjuncts the factor mentions + | core + core ::= atom # array, getitem, delta(idx, w), ... + | R.reduce(prod, ranges) # every factor in prod uses a range var + cond ::= And.plus(cmp, ...) # comparison atoms, stream-var first + ranges ::= { var: range(0, N, 1), ... } + +The invariants, grouped by the rules that establish them: + +A. Sums are factored over products (the einsum shape). + - :class:`Factor` pulls stream-invariant product factors into outer product + factors. :class:`ReduceSplit` distributes a commutative reduction over a + same-monoid ``plus``; each resulting reduction initially retains the full + stream mapping, after which factorization and stream-elimination rules may + simplify streams unused by an individual summand. + - Explicit (materialized) products-of-sums are expanded to sums-of-products + by :class:`PlusDistr`. (This is the opposite direction from factorization, + but it acts on explicit ``plus`` terms, never on ``reduce`` nodes, so the + two never fight.) + +B. Streams are independent ranges ``range(0, N, 1)``. + Every non-range binding has a rule that removes it: eager arrays + (:class:`~effectful.handlers.jax.monoid.ReduceArrayGather` + + :class:`EliminateSingletonStreams`), unions (:class:`ReduceUnion`), weighted + streams (:class:`ReduceWeightedStream`), cartesian products (eliminated by + inversion in :class:`ReduceDistributeCartesianProduct`), dependent ranges + (:class:`~effectful.handlers.jax.monoid.ReduceDependentRangeMask`), and + symbolic length-1 streams (:class:`EliminateSingletonStreams`). In particular + ``CartesianProduct.reduce`` is fully eliminated. + +C. ``plus``, ``mask``, and ``delta`` stay symbolic (not lowered). + - ``plus`` is only normalized structurally: flattened/associative + (:class:`PlusAssoc`), with nullary/unary collapsed + (:class:`PlusEmpty`/:class:`PlusSingle`) and int/float operands unified + (:class:`PlusCastFloat`). The scalar/array implementations live in + ``EvaluateIntp``, not here. + - ``mask`` floats to a canonical position: fused to a single layer with a + conjunctive (``And.plus``) condition (:class:`MaskFusion`), constant-bool + conditions discharged (:class:`MaskBool`), pushed *down* onto the factors + of a ``plus`` (:class:`MaskPushPlus`) -- each conjunct landing on the + factors that mention its variables, so a mask sits adjacent to a single + product factor where it can fuse with a gather (a conjunct no factor + mentions stays as a residual outer mask) -- and out of a ``reduce`` when + the condition is stream-independent (:class:`ReduceMaskHoist`). A mask + remaining inside a reduce body therefore depends on a reduced stream. + Comparison atoms are oriented stream-variable-first + (:class:`MaskOrderStreamOps`). + - ``delta`` is flattened: nested deltas merge (:class:`DeltaFusion`), + empty-index deltas collapse to their weight (:class:`DeltaEmpty`), and + subscripting a delta becomes a mask (:class:`GetitemDelta`). A normal-form + delta has a non-empty index, a single layer, and is never subscripted. + +D. Structural invariants. + - No same-monoid ``reduce`` is nested directly in a ``reduce`` body: they are + fused into one reduce over the combined stream set (:class:`ReduceFusion`), + so each reduce is maximal in its stream set per monoid. + - Containers are outermost: :class:`MonoidOverCallable`, + :class:`MonoidOverMapping`, and :class:`MonoidOverSequence` push + ``reduce``/``plus`` through lambdas, mappings, and sequences down to the + scalar/array leaves, so a monoid op never wraps a container. + +Termination and confluence +=========================== + +The rules run as a deterministic, priority-ordered, innermost normalization +(``evaluate`` rewrites args before parents, and each constructed replacement +re-enters the interpretation to a fixpoint). For a fixed *syntactic* input the +result is therefore unique. + +The normal form is **not fully canonical**, however: semantically equal but +differently-presented inputs can reach distinct (still semantically equal) +forms. Commutative ``plus`` arguments are sorted by :class:`PlusOrder`, and +consecutive syntactically equal arguments are removed for idempotent monoids by +:class:`PlusConsecutiveDups`. This provides limited AC normalization, but the +sort key uses object identity to break syntactic-hash ties, and duplicates are +preserved for non-idempotent monoids. Moreover, variable-elimination order +depends on ``streams`` insertion order via +:func:`choose_contraction`/:func:`outer_stream`, while argument order remains +significant for noncommutative monoids. Downstream code therefore relies only +on the *semantic* normal form (the result is ultimately evaluated to a concrete +array), not on full syntactic canonicity. + +Termination rests on per-family progress (each "lowering" rule strictly +consumes a resource -- a cartesian/weighted/union/array/singleton stream, a +nested reduce/mask/delta, or a liftable factor) plus redex-shape disjointness +that keeps the families from cycling (e.g. :class:`ReduceUnfactor` is kept out +of ``NormalizeIntp`` because it would form a rewrite cycle with +:class:`Factor`). There +is no single global measure and no confluence theorem: the system is +normalizing by construction, so a new rule that overlaps an existing redex shape +can introduce a loop or shift which normal form is reached. Note also that +:class:`PlusDistr` is exponential in the number of distributed sums. """ diff --git a/effectful/ops/syntax.py b/effectful/ops/syntax.py index 5d0b1e983..b5ef5bbba 100644 --- a/effectful/ops/syntax.py +++ b/effectful/ops/syntax.py @@ -5,7 +5,7 @@ import numbers import operator import typing -from collections.abc import Callable, Iterable, Mapping +from collections.abc import Callable, Iterable, KeysView, Mapping, ValuesView from typing import Annotated, Any from effectful.ops.types import ( @@ -606,11 +606,6 @@ def __init__( self._args = args self._kwargs = kwargs - def __eq__(self, other) -> bool: - from effectful.ops.syntax import syntactic_eq - - return syntactic_eq(self, other) - @property def op(self): return self._op @@ -790,6 +785,111 @@ def __next__(self: collections.abc.Iterator[T]) -> T: next_ = _IteratorTerm.__next__ +@defdata.register(collections.abc.Collection) +class _CollectionTerm[T](_IterableTerm[T], collections.abc.Collection[T]): + @defop + def __contains__(self: collections.abc.Collection[T], x: T) -> bool: + if not isinstance(self, Term) and not isinstance(x, Term): + return x in self + else: + raise NotHandled + + @defop + def __len__(self: collections.abc.Collection[T]) -> int: + if not isinstance(self, Term): + return len(self) + else: + raise NotHandled + + +@defdata.register(collections.abc.Sequence) +class _SequenceTerm[T](_CollectionTerm[T], collections.abc.Sequence[T]): + @Operation.define + def __getitem__(self: collections.abc.Sequence[T], key: int) -> T: + if not isinstance(self, Term) and not isinstance(key, Term): + return self[key] + else: + raise NotHandled + + @Operation.define + def __reversed__(self: collections.abc.Sequence[T]) -> collections.abc.Iterator[T]: + if not isinstance(self, Term): + return reversed(self) + else: + raise NotHandled + + @Operation.define + def index(self: collections.abc.Sequence[T], *args, **kwargs) -> int: + from effectful.ops.semantics import fvsof + + if not fvsof((self, *args, *kwargs.values())): + return self.index(*args, **kwargs) + else: + raise NotHandled + + @Operation.define + def count(self: collections.abc.Sequence[T], value: T) -> int: + from effectful.ops.semantics import fvsof + + if not fvsof((self, value)): + return self.count(value) + else: + raise NotHandled + + +@defdata.register(collections.abc.Mapping) +class _MappingTerm[K, V](_CollectionTerm[K]): + @defop + def __getitem__(self: collections.abc.Mapping[K, V], key: K) -> V: + from effectful.ops.semantics import fvsof + + if not isinstance(self, Term) and not fvsof(key): + return self[key] + else: + raise NotHandled + + @defop + def get( + self: collections.abc.Mapping[K, V], key: K, default: V | None = None + ) -> V | None: + from effectful.ops.semantics import fvsof + + if not isinstance(self, Term) and not fvsof(key): + return self.get(key, default) + else: + raise NotHandled + + @defop + def keys(self: collections.abc.Mapping[K, V]) -> KeysView[K]: + if not isinstance(self, Term): + return self.keys() + else: + raise NotHandled + + @defop + def values(self: collections.abc.Mapping[K, V]) -> ValuesView[V]: + if not isinstance(self, Term): + return self.values() + else: + raise NotHandled + + @defop + def __eq__(self: collections.abc.Mapping[K, V], other) -> bool: + if not isinstance(self, Term) and not isinstance(other, Term): + return self == other + else: + raise NotHandled + + +@Operation.define +def as_dict[K, V](*args: tuple[K, V]) -> Mapping[K, V]: + from effectful.ops.semantics import fvsof + + if not fvsof(args): + return dict(args) + raise NotHandled + + @_CustomSingleDispatchCallable def syntactic_eq( __dispatch: Callable[[type], Callable[[Any, Any], bool]], x, other @@ -891,7 +991,8 @@ def _(x: collections.abc.Sequence, other) -> bool: ) else: return ( - isinstance(other, collections.abc.Sequence) + not isinstance(other, Term) + and isinstance(other, collections.abc.Sequence) and len(x) == len(other) and all(syntactic_eq(a, b) for a, b in zip(x, other)) ) @@ -1246,36 +1347,61 @@ def __index__(self) -> int: def __eq__(self, other) -> bool: if not isinstance(self, Term) and not isinstance(other, Term): return self.__eq__(other) - else: - return syntactic_eq(self, other) + + if not isinstance(other, numbers.Number): + return NotImplemented + + raise NotHandled + + @defop + def __ne__(self, other) -> bool: + if not isinstance(self, Term) and not isinstance(other, Term): + return self.__ne__(other) + + if not isinstance(other, numbers.Number): + return NotImplemented + + raise NotHandled @defop def __lt__(self, other) -> bool: if not isinstance(self, Term) and not isinstance(other, Term): return self.__lt__(other) - else: - raise NotHandled + + if not isinstance(other, numbers.Number): + return NotImplemented + + raise NotHandled @defop def __gt__(self, other) -> bool: if not isinstance(self, Term) and not isinstance(other, Term): return self.__gt__(other) - else: - raise NotHandled + + if not isinstance(other, numbers.Number): + return NotImplemented + + raise NotHandled @defop def __le__(self, other) -> bool: if not isinstance(self, Term) and not isinstance(other, Term): return self.__le__(other) - else: - raise NotHandled + + if not isinstance(other, numbers.Number): + return NotImplemented + + raise NotHandled @defop def __ge__(self, other) -> bool: - if not isinstance(self, Term) and not isinstance(other, Term): + if not isinstance(self, _NumberTerm) and not isinstance(other, _NumberTerm): return self.__ge__(other) - else: - raise NotHandled + + if not isinstance(other, numbers.Number): + return NotImplemented + + raise NotHandled @defop def __add__(self, other: T) -> T: @@ -1487,6 +1613,53 @@ class _BoolTerm[T: bool](_IntegralTerm[T]): # type: ignore pass +@Operation.define +def ite[T](cond, then: T, else_: T) -> T: + """If-then-else operation.""" + # Note: cond is specifically not annotated to allow this operation to take + # non-bool arguments (like boolean arrays). + if not isinstance(cond, Term): + return then if cond else else_ + raise NotHandled + + +@Operation.define +def range_(stop: int, *args: int) -> range: + if any(isinstance(x, Term) for x in (stop, *args)): + raise NotHandled + return range(stop, *args) + + +@defdata.register(range) +class _RangeTerm(_SequenceTerm[int]): + @property + @Operation.define + def start(self) -> int: + if not isinstance(self, Term): + return self.start + if self.op == range_: + return 0 if len(self.args) < 2 else self.args[0] + raise NotHandled + + @property + @Operation.define + def stop(self) -> int: + if not isinstance(self, Term): + return self.stop + if self.op == range_: + return self.args[0] if len(self.args) < 2 else self.args[1] + raise NotHandled + + @property + @Operation.define + def step(self) -> int: + if not isinstance(self, Term): + return self.step + if self.op == range_: + return 1 if len(self.args) < 3 else self.args[2] + raise NotHandled + + class ConstructorOperation[**Q, V](Operation[Q, V]): @classmethod @functools.cache diff --git a/effectful/ops/types.py b/effectful/ops/types.py index 018932eb6..6f119e2fa 100644 --- a/effectful/ops/types.py +++ b/effectful/ops/types.py @@ -279,11 +279,17 @@ def _define_callable[**P, T]( cls, t: Callable[P, T], *, name: str | None = None ) -> "Operation[P, T]": if isinstance(t, Operation): + sig = inspect.signature(t) @functools.wraps(t) def func(*args, **kwargs): raise NotHandled + # functools.wraps does not copy the signature. Instead it points to + # the wrapped function. inspect.signature will traverse this chain + # unless it is removed. + del func.__wrapped__ + func.__signature__ = sig # type: ignore[attr-defined] op = cls.define(func, name=name) else: op = cls(t, name=name) # type: ignore[arg-type] @@ -505,7 +511,14 @@ def _instance_op(instance, *args, **kwargs): else: return default_result - name = ("" if owner is None else f"{owner.__name__}_") + self.__name__ + name: str = "" + if instance is not None and hasattr(instance, "__name__"): + assert isinstance(instance.__name__, str) + name = instance.__name__ + elif owner is not None: + name = owner.__name__ + name += f"_{self.__name__}" + instance_op = self.define( types.MethodType(_instance_op, instance), name=name ) diff --git a/pyproject.toml b/pyproject.toml index 29dca1eb3..f1d4a672d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,16 +48,14 @@ numpyro = [ "jax<0.10" ] llm = [ - # 1.92.0 ships no wheel, so it builds from sdist; on Python 3.14 that build pulls - # pyo3-ffi 0.23.5, which doesn't support 3.14, so `uv sync` fails the whole matrix. - "litellm!=1.92.0", + "litellm<1.92", "tenacity", "mypy", "autoflake", "pillow", "pydantic", "typing_extensions", - "restrictedpython>=8.1" + "restrictedpython>=8.1", ] prettyprinter = ["prettyprinter"] docs = [ diff --git a/tests/_monoid_helpers.py b/tests/_monoid_helpers.py index 72787558a..88b3b436d 100644 --- a/tests/_monoid_helpers.py +++ b/tests/_monoid_helpers.py @@ -12,8 +12,13 @@ import effectful.handlers.jax.numpy as _jnp from effectful.internals.runtime import interpreter -from effectful.ops.monoid import NormalizeIntp, Stream, _is_monoid_weighted -from effectful.ops.semantics import apply, evaluate, fvsof, handler +from effectful.ops.monoid import ( + EvaluateIntp, + NormalizeIntp, + Stream, + _is_monoid_weighted, +) +from effectful.ops.semantics import apply, coproduct, evaluate, fvsof, handler from effectful.ops.syntax import _BaseTerm, defdata, deffn, syntactic_eq from effectful.ops.types import NotHandled, Operation, Term @@ -40,9 +45,9 @@ def _canonical_op(idx: int, op: Operation) -> Operation: if idx in _op_cache: return _op_cache[idx] - op = Operation.define(op, name=f"__cv_{idx}") - _op_cache[idx] = op - return op + canon_op = Operation.define(op, name=f"__cv_{idx}") + _op_cache[idx] = canon_op + return canon_op cx = _canonicalize(x, _canonical_op) cy = _canonicalize(y, _canonical_op) @@ -195,15 +200,14 @@ def define_vars(self, *names: str, **kwargs) -> Operation | tuple[Operation, ... return tuple(self._fresh_op(n, **kwargs) for n in names) def check_rewrite( - self, - lhs, - rhs, - rule, - *, - max_examples: int = 25, - deadline=None, - normalize=NormalizeIntp, + self, lhs, rhs, rule, *, max_examples: int = 25, deadline=None, normalize=None ) -> None: + normalize = ( + normalize + if normalize is not None + else coproduct(EvaluateIntp, NormalizeIntp) + ) + with handler(rule): norm = evaluate(lhs) assert syntactic_eq_alpha(norm, rhs) diff --git a/tests/test_handlers_jax.py b/tests/test_handlers_jax.py index e07385725..4c825542c 100644 --- a/tests/test_handlers_jax.py +++ b/tests/test_handlers_jax.py @@ -2,7 +2,7 @@ import pytest import effectful.handlers.jax.numpy as jnp -from effectful.handlers.jax import bind_dims, jax_getitem, jit, sizesof +from effectful.handlers.jax import bind_dims, jax_getitem, jit, sizesof, unbind_dims from effectful.ops.semantics import evaluate, fvsof, handler from effectful.ops.syntax import defdata, defop, syntactic_eq from effectful.ops.types import Term @@ -80,6 +80,9 @@ def test_bind_dims(): t9 = evaluate(t8) assert not (fvsof(t9) & {i, j, k, w}) + t10 = bind_dims(jax_getitem(jnp.ones((5, 6)), [unbind_dims(w(), i), i()]), i) + assert isinstance(t10, Term) and t10.op == bind_dims + def test_tpe_1(): i, j = defop(jax.Array), defop(jax.Array) diff --git a/tests/test_handlers_jax_monoid.py b/tests/test_handlers_jax_monoid.py index e410a3e69..14a1db08e 100644 --- a/tests/test_handlers_jax_monoid.py +++ b/tests/test_handlers_jax_monoid.py @@ -1,33 +1,35 @@ -import functools +import sys import jax +import jax.dlpack +import numpy as np +import pyro.ops.contract import pytest +import torch from jax import random as random import effectful.handlers.jax.numpy as jnp from effectful.handlers.jax import bind_dims, jax_getitem, unbind_dims from effectful.handlers.jax.monoid import ( ARRAY_REDUCTORS, - DeltaEmpty, ReduceArray, ReduceArrayGather, ReduceDeltaSimpleRange, - ReduceDependentRangeMask, ReduceSumProductContraction, - delta, einsum, ) from effectful.ops.monoid import ( EliminateSingletonStreams, + EvaluateIntp, NormalizeIntp, Product, Sum, ) -from effectful.ops.semantics import coproduct, handler +from effectful.ops.semantics import coproduct, evaluate, handler from tests._monoid_helpers import JaxBackend MONOIDS = [ - pytest.param(monoid, reductor, id=monoid._name) + pytest.param(monoid, reductor, id=monoid.__name__) for (monoid, reductor) in ARRAY_REDUCTORS.items() ] @@ -94,105 +96,32 @@ def test_reduce_array_gather_dep(monoid, reductor, backend: JaxBackend): @pytest.mark.parametrize("monoid,reductor", MONOIDS) def test_reduce_array_1(monoid, reductor, backend: JaxBackend): - (x, k) = backend.define_vars("x", "k", ret="scalar") + x = backend.define_vars("x", ret="scalar") + arr = jnp.arange(4, 9) X = jnp.arange(5) - lhs = monoid.reduce(x(), {x: X}) - rhs = reductor(bind_dims(unbind_dims(X, k), k), axis=(0,)) - backend.check_rewrite( - lhs=lhs, - rhs=rhs, - rule=functools.reduce( - coproduct, # type: ignore[arg-type] - [ - ReduceArrayGather(), - EliminateSingletonStreams(), - ReduceArray(), - ReduceDeltaSimpleRange(), - ], - ), - ) + with handler(NormalizeIntp), handler(EvaluateIntp): + actual = monoid.reduce(unbind_dims(arr, x), {x: X}) + assert jnp.allclose(actual, reductor(arr)) @pytest.mark.parametrize("monoid,reductor", MONOIDS) def test_reduce_array_2(monoid, reductor, backend: JaxBackend): - (x, y, k1, k2) = backend.define_vars("x", "y", "k1", "k2", ret="scalar") + (x, y) = backend.define_vars("x", "y", ret="scalar") X = jnp.arange(5) Y = jnp.arange(7) - f = backend.define_vars( - "f", arg_types=(backend.scalar_typ, backend.scalar_typ), ret="scalar" - ) - - lhs = monoid.reduce(f(x(), y()), {x: X, y: Y}) - rhs = reductor( - bind_dims(f(unbind_dims(X, k1), unbind_dims(Y, k2)), k1, k2), axis=(0, 1) - ) - backend.check_rewrite( - lhs=lhs, - rhs=rhs, - rule=functools.reduce( - coproduct, # type: ignore[arg-type] - [ - ReduceArrayGather(), - EliminateSingletonStreams(), - ReduceArray(), - ReduceDeltaSimpleRange(), - ], - ), - ) - - -@pytest.mark.parametrize("monoid,reductor", MONOIDS) -def test_reduce_array_3(monoid, reductor, backend: JaxBackend): - """Stream `y` is `g(x())` — depends on the bound element of X. The reducer - must inline ``g`` along the same named dim used to unbind `x`.""" - (x, y, k1, k2) = backend.define_vars("x", "y", "k1", "k2", ret="scalar") - X = jnp.arange(5) + arr_X = jnp.arange(4, 9) + arr_Y = jnp.arange(4, 11) - f = backend.define_vars( - "f", arg_types=[backend.scalar_typ, backend.scalar_typ], ret="scalar" - ) - g = backend.define_vars("g", arg_types=[backend.scalar_typ], ret="stream") - - lhs = monoid.reduce(f(x(), y()), {x: X, y: g(x())}) - rhs = reductor( - bind_dims( - monoid.reduce(f(unbind_dims(X, x), y()), {y: g(unbind_dims(X, x))}), x - ), - axis=(0,), - ) - backend.check_rewrite( - lhs=lhs, - rhs=rhs, - rule=functools.reduce( - coproduct, # type: ignore[arg-type] - [ - ReduceArrayGather(), - EliminateSingletonStreams(), - ReduceArray(), - ReduceDeltaSimpleRange(), - DeltaEmpty(), - ], - ), - ) - - -@pytest.mark.parametrize("monoid,reductor", MONOIDS) -def test_arange_reduce_direct_full(monoid, reductor, backend: JaxBackend): - """A full-range direct index ``A[v()]`` over ``v: arange(N)`` slices the - whole axis (``A[0:N:1]``) and reduces it -- no materialized-arange gather. - """ - (v, k) = backend.define_vars("v", "k", ret="scalar") - A = backend.define_vars("A", ret="stream") + with handler(NormalizeIntp), handler(EvaluateIntp): + actual = monoid.reduce( + monoid.plus(unbind_dims(arr_X, x), unbind_dims(arr_Y, y)), {x: X, y: Y} + ) + expected = reductor( + bind_dims(monoid.plus(unbind_dims(arr_X, x), unbind_dims(arr_Y, y)), x, y) + ) - lhs = monoid.reduce(jax_getitem(A(), [v()]), {v: range(7)}) - rhs = reductor( - bind_dims(jax_getitem(jax_getitem(A(), [slice(0, 7, 1)]), [k()]), k), - axis=(0,), - ) - backend.check_rewrite( - lhs=lhs, rhs=rhs, rule=coproduct(ReduceArray(), ReduceDeltaSimpleRange()) - ) + assert jnp.allclose(actual, expected) @pytest.mark.parametrize("monoid,reductor", MONOIDS) @@ -244,28 +173,12 @@ def test_arange_reduce_two_streams(monoid, reductor, backend: JaxBackend): # --------------------------------------------------------------------------- -@pytest.mark.parametrize("monoid,reductor", MONOIDS) -def test_reduce_delta_empty(monoid, reductor, backend: JaxBackend): - """An empty-index delta unwraps to its body. - - reduce(M, streams, delta((), body)) ≡ reduce(M, streams, body) - """ - x = backend.define_vars("x", ret="scalar") - X = backend.define_vars("X", ret="stream") - - lhs = monoid.reduce(delta((), x()), {x: X()}) - rhs = monoid.reduce(x(), {x: X()}) - backend.check_rewrite( - lhs=lhs, rhs=rhs, rule=coproduct(ReduceDeltaSimpleRange(), DeltaEmpty()) - ) - - @pytest.mark.parametrize("monoid,reductor", MONOIDS) def test_reduce_delta_empty_arange(monoid, reductor, backend: JaxBackend): x = backend.define_vars("x", ret="scalar") f = backend.define_vars("f", arg_types=[backend.scalar_typ], ret="scalar") - lhs = monoid.reduce(delta((x(),), f(x())), {x: range(0)}) + lhs = monoid.reduce(monoid.delta((x(),), f(x())), {x: range(0)}) rhs = bind_dims(f(unbind_dims(jnp.array([]), x)), x) backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceDeltaSimpleRange()) @@ -282,7 +195,7 @@ def test_reduce_delta_independent_one(monoid, reductor, backend: JaxBackend): # We use a concrete range here instead of an abstract one, because # unbind_dims is undefined on empty arrays (and the rewrite produces a # different rhs in this case) - lhs = monoid.reduce(delta((y(),), f(y())), {y: range(3)}) + lhs = monoid.reduce(monoid.delta((y(),), f(y())), {y: range(3)}) rhs = bind_dims(f(unbind_dims(jnp.arange(3), k)), k) backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceDeltaSimpleRange()) @@ -302,7 +215,9 @@ def test_reduce_delta_independent_preserves_others( "f", arg_types=[backend.scalar_typ, backend.scalar_typ], ret="scalar" ) - lhs = monoid.reduce(delta((x(), y()), f(x(), y())), {x: range(2), y: range(3)}) + lhs = monoid.reduce( + monoid.delta((x(), y()), f(x(), y())), {x: range(2), y: range(3)} + ) rhs = bind_dims( bind_dims(f(unbind_dims(jnp.arange(2), x), unbind_dims(jnp.arange(3), k)), k), x ) @@ -315,12 +230,12 @@ def test_reduce_delta_simple_dep(monoid, reductor, backend: JaxBackend): X = jnp.arange(3) lhs = monoid.reduce( - delta((x(),), unbind_dims(X, x) + y()), + monoid.delta((x(),), unbind_dims(X, x) + y()), {x: range(3), y: jnp.stack([x(), x() + 1])}, ) rhs = bind_dims( monoid.reduce( - delta((), unbind_dims(X, x) + y()), + unbind_dims(X, x) + y(), { y: jnp.stack( [unbind_dims(jnp.arange(3), x), unbind_dims(jnp.arange(3), x) + 1] @@ -332,55 +247,6 @@ def test_reduce_delta_simple_dep(monoid, reductor, backend: JaxBackend): backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceDeltaSimpleRange()) -@pytest.mark.parametrize("monoid,reductor", MONOIDS) -def test_reduce_dependent_range_mask(monoid, reductor, backend: JaxBackend): - """A dependent range stream gets rewritten to the referent's bbox stream, - with the original constraint folded into the body as a where-guard. - - reduce(M, {u: range(0, N, 1), v: range(0, u(), 1)}, body) - ≡ reduce(M, {u: range(0, N, 1), v: range(0, N, 1)}, where(v() < u(), body, M.identity)) - """ - (u, v) = backend.define_vars("u", "v", ret="scalar") - N = 5 - f = backend.define_vars( - "f", arg_types=[backend.scalar_typ, backend.scalar_typ], ret="scalar" - ) - - body = f(u(), v()) - - lhs = monoid.reduce(body, {u: range(N), v: jnp.arange(u())}) - rhs = monoid.reduce( - jnp.where(v() < u(), body, monoid.identity), {u: range(N), v: range(N)} - ) - backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceDependentRangeMask()) - - -@pytest.mark.parametrize("monoid,reductor", MONOIDS) -def test_reduce_dependent_range_mask_delta_body(monoid, reductor, backend: JaxBackend): - """When the body is a delta term, R4 folds the constraint into the delta's - weight while leaving its index tuple untouched. - - reduce(M, {u: range(N), v: range(u())}, delta((u(), v()), w)) - ≡ reduce(M, {u: range(N), v: range(N)}, - delta((u(), v()), where(v() < u(), w, M.identity))) - """ - (u, v) = backend.define_vars("u", "v", ret="scalar") - N = 5 - f = backend.define_vars( - "f", arg_types=[backend.scalar_typ, backend.scalar_typ], ret="scalar" - ) - - weight = f(u(), v()) - idx = (u(), v()) - - lhs = monoid.reduce(delta(idx, weight), {u: range(N), v: jnp.arange(u())}) - rhs = monoid.reduce( - delta(idx, jnp.where(v() < u(), weight, monoid.identity)), - {u: range(N), v: range(N)}, - ) - backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceDependentRangeMask()) - - def test_reduce_contraction_single(backend: JaxBackend): i = backend.define_vars("i", ret="scalar") (A, B) = backend.define_vars( @@ -390,8 +256,8 @@ def test_reduce_contraction_single(backend: JaxBackend): lhs = Sum.reduce(Product.plus(A(i()), B(i())), {i: range(5)}) rhs = jnp.einsum( "a...,a...->...", - Sum.reduce(delta((i(),), A(i())), {i: range(5)}), - Sum.reduce(delta((i(),), B(i())), {i: range(5)}), + Sum.reduce(Sum.delta((i(),), A(i())), {i: range(5)}), + Sum.reduce(Sum.delta((i(),), B(i())), {i: range(5)}), ) backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceSumProductContraction()) @@ -405,8 +271,8 @@ def test_reduce_contraction_double(backend: JaxBackend): lhs = Sum.reduce(Product.plus(A(i(), j()), B(i(), j())), {i: range(5), j: range(7)}) rhs = jnp.einsum( "ab...,ab...->...", - Sum.reduce(delta((i(), j()), A(i(), j())), {i: range(5), j: range(7)}), - Sum.reduce(delta((i(), j()), B(i(), j())), {i: range(5), j: range(7)}), + Sum.reduce(Sum.delta((i(), j()), A(i(), j())), {i: range(5), j: range(7)}), + Sum.reduce(Sum.delta((i(), j()), B(i(), j())), {i: range(5), j: range(7)}), ) backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceSumProductContraction()) @@ -422,11 +288,16 @@ def test_reduce_matmul(backend: JaxBackend): (b, i, j, k) = backend.define_vars("b", "i", "j", "k", ret="scalar") with handler(NormalizeIntp): - actual = Sum.reduce( - delta((b(), i(), k()), unbind_dims(X, b, i, j) * unbind_dims(Y, b, j, k)), + norm = Sum.reduce( + Sum.delta( + (b(), i(), k()), unbind_dims(X, b, i, j) * unbind_dims(Y, b, j, k) + ), {b: range(B), i: range(I), j: range(J), k: range(K)}, ) + with handler(EvaluateIntp), handler(NormalizeIntp): + actual = evaluate(norm) + assert isinstance(actual, jax.Array) expected = jnp.einsum("bij,bjk->bik", X, Y) assert jnp.allclose(actual, expected) @@ -553,6 +424,63 @@ def test_einsum_matches_jnp(spec: str, sizes, rng_key): assert actual.shape == expected.shape, ( f"shape mismatch for {spec!r}: got {actual.shape}, expected {expected.shape}" ) - assert jnp.allclose(actual, expected, atol=1e-4, rtol=1e-4), ( + assert jnp.allclose(actual, expected, atol=1e-4, rtol=1e-3), ( f"value mismatch for {spec!r}" ) + + +# see https://github.com/pyro-ppl/pyro/blob/dev/tests/ops/test_contract.py +# Let abcde be enum dims and ijk be plates. +PLATED_EINSUM_CASES = [ + ("abi,abi->", "i"), + ("abi,b->", "i"), + ("abi,b->i", "i"), + ("abi,b->b", "i"), + ("abi,b->ai", "i"), + ("aij,bi->", "ij"), + ("acij,bi->", "ij"), + ("aij,bi,c->", "ij"), + ("abij,bi->aij", "ij"), + ("abij,bci->acij", "ij"), + ("abij,bci->", "ij"), + ("ab,bcij->", "ij"), + ("ija,ika->", "ijk"), + ("ab,bcdi,defij,fgijk->", "ijk"), +] + + +@pytest.mark.parametrize("spec,plates", PLATED_EINSUM_CASES) +def test_plated_einsum(spec, plates, rng_key): + sys.setrecursionlimit(10_000) + + def _to_torch(arr): + return torch.from_dlpack(arr) + + def _to_mapping(arr): + return {i: float(arr[*i]) for i in np.ndindex(arr.shape)} + + dim_sizes = { + "a": 2, + "b": 3, + "c": 4, + "d": 5, + "e": 6, + "f": 7, + "g": 8, + "i": 2, + "j": 3, + "k": 4, + } + + operands = _make_operands(spec, dim_sizes, rng_key) + torch_operands = (_to_torch(op) for op in operands) + + try: + expected = pyro.ops.contract.naive_ubersum( + spec, *torch_operands, plates=plates, backend="torch" + )[0] + except NotImplementedError: + pytest.skip("Not implemented by pyro.ops.contract.einsum") + + actual = einsum(spec, *operands, plates=plates) + assert torch.allclose(torch.tensor(actual), expected, atol=1e-4, rtol=1e-4) diff --git a/tests/test_ops_monoid.py b/tests/test_ops_monoid.py index 9976fd6dc..41fc01aa4 100644 --- a/tests/test_ops_monoid.py +++ b/tests/test_ops_monoid.py @@ -1,44 +1,60 @@ import math +import sys import typing -from collections.abc import Iterable +from collections.abc import Iterable, Mapping import pytest from hypothesis import HealthCheck, given, settings from hypothesis import strategies as st import effectful.handlers.jax.monoid # noqa: F401 -import effectful.handlers.jax.numpy as jnp from effectful.ops.monoid import ( + And, CartesianProduct, + CartesianProductPlus, EliminateSingletonStreams, + EvaluateIntp, + Factor, Max, Min, Monoid, MonoidOverMapping, MonoidOverSequence, NormalizeIntp, + Or, PlusAssoc, PlusConsecutiveDups, PlusDistr, - PlusDups, PlusEmpty, + PlusOrder, PlusSingle, Product, - ReduceCartesianWeightedStream, + ReduceDependentRangeMask, + ReduceDisjunctiveDisequalityMask, ReduceDistributeCartesianProduct, - ReduceFactorization, + ReduceEqualityMaskRange, ReduceFusion, + ReduceMaskHoist, ReducePartial, ReduceSplit, + ReduceUnfactor, + ReduceUnion, ReduceWeightedStream, + ReduceWhereEqualityPeel, + ReduceWhereToMasks, Sum, + Union, + WhereHoist, distributes_over, + is_commutative, ) from effectful.ops.semantics import coproduct, evaluate, fvsof, handler -from effectful.ops.syntax import deffn +from effectful.ops.syntax import as_dict, ite, range_, syntactic_eq from effectful.ops.types import NotHandled, Operation, Term from tests._monoid_helpers import Backend, IntBackend, JaxBackend, syntactic_eq_alpha +sys.setrecursionlimit(10_000) + @pytest.fixture(params=[IntBackend, JaxBackend], ids=["int", "jax"]) def backend(request) -> Backend: @@ -80,6 +96,127 @@ def backend(request) -> Backend: ) ] +COMMUTATIVE_MONOID_PAIRS = [ + v + for v in MONOID_PAIRS + if all(is_commutative(typing.cast(Monoid, vv)) for vv in v.values) +] + + +def test_plus_ite_hoist(backend: Backend): + """Monoid addition distributes into both branches of an ``ite``.""" + cond_lhs, cond_rhs, a, b, c = backend.define_vars( + "cond_lhs", "cond_rhs", "a", "b", "c", ret="scalar" + ) + cond = cond_lhs() == cond_rhs() + lhs = Product.plus(ite(cond, a(), b()), c()) + rhs = ite(cond, Product.plus(a(), c()), Product.plus(b(), c())) + + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=WhereHoist()) + + +def test_reduce_ite_hoist(backend: Backend): + """An ``ite`` with a stream-independent condition hoists.""" + i, out_i, out_j = backend.define_vars("i", "out_i", "out_j", ret="scalar") + f, g = backend.define_vars("f", "g", arg_types=(backend.scalar_typ,), ret="scalar") + streams = {i: range(3)} + cond = out_i() == out_j() + + lhs = Sum.reduce(ite(cond, f(i()), g(i())), streams) + rhs = ite(cond, Sum.reduce(f(i()), streams), Sum.reduce(g(i()), streams)) + + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=WhereHoist()) + + +def test_reduce_ite_hoist_dependent_noop(backend: Backend): + """An ``ite`` whose condition uses the reduced stream remains in place.""" + i = backend.define_vars("i", ret="scalar") + f, g = backend.define_vars("f", "g", arg_types=(backend.scalar_typ,), ret="scalar") + term = Sum.reduce(ite(i() == 0, f(i()), g(i())), {i: range(3)}) + + backend.check_rewrite(lhs=term, rhs=term, rule=WhereHoist()) + + +def test_reduce_ite_equality_peel(backend: Backend): + """An independent conjunct peels off a stream-selecting equality guard.""" + i, j = backend.define_vars("i", "j", ret="scalar") + f, g = backend.define_vars("f", "g", arg_types=(backend.scalar_typ,), ret="scalar") + streams = {j: range(3)} + outer = 0 == i() + inner = 0 == j() + + lhs = Product.reduce(ite(And.plus(inner, outer), f(j()), g(j())), streams) + rhs = ite( + And.plus(outer), + Product.reduce(ite(And.plus(inner), f(j()), g(j())), streams), + Product.reduce(g(j()), streams), + ) + + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceWhereEqualityPeel()) + + +def test_reduce_ite_equality_peel_requires_stream_equality(backend: Backend): + """A mixed guard without a stream equality is left unchanged.""" + j, out_i, i = backend.define_vars("j", "out_i", "i", ret="scalar") + f, g = backend.define_vars("f", "g", arg_types=(backend.scalar_typ,), ret="scalar") + term = Product.reduce( + ite(And.plus(j() < 2, out_i() == i()), f(j()), g(j())), + {j: range(3)}, + ) + + backend.check_rewrite(lhs=term, rhs=term, rule=ReduceWhereEqualityPeel()) + + +def test_reduce_disjunctive_disequality_mask(): + backend = IntBackend() + i, j, out_i, out_j, value = backend.define_vars( + "i", "j", "out_i", "out_j", "value", ret="scalar" + ) + streams = {i: range(3), j: range(4)} + + lhs = Sum.reduce( + Sum.mask(value(), Or.plus(i() != out_i(), j() != out_j())), streams + ) + rhs = Sum.plus( + Sum.reduce(Sum.mask(value(), And.plus(i() != out_i())), streams), + Sum.reduce( + Sum.mask(value(), And.plus(i() == out_i(), j() != out_j())), streams + ), + ) + + with handler(ReduceDisjunctiveDisequalityMask()): + actual = evaluate(lhs) + assert syntactic_eq_alpha(actual, rhs) + + +def test_reduce_where_to_masks(): + backend = IntBackend() + """A conjunctive equality where partitions into complementary masks.""" + i, j = backend.define_vars("i", "j", ret="scalar") + f, g = backend.define_vars("f", "g", arg_types=(backend.scalar_typ,), ret="scalar") + streams = {i: range(2), j: range(3)} + + lhs = Product.reduce(ite(And.plus(i() == 0, j() == 0), f(i()), g(j())), streams) + rhs = Product.plus( + Product.reduce(Product.mask(f(i()), And.plus(i() == 0, j() == 0)), streams), + Product.reduce(Product.mask(g(j()), Or.plus(i() != 0, j() != 0)), streams), + ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceWhereToMasks()) + + +def test_reduce_where_to_masks_requires_stream_equalities(backend: JaxBackend): + """Independent equality conjuncts are left for WhereEqualityPeel.""" + i, out_i, x, y = backend.define_vars("i", "out_i", "x", "y", ret="scalar") + f, g = backend.define_vars("f", "g", arg_types=(backend.scalar_typ,), ret="scalar") + term = Product.reduce( + ite(And.plus(i() == out_i(), x() == y()), f(i()), g(i())), + {i: range(2)}, + ) + + with handler(ReduceWhereToMasks()): + actual = evaluate(term) + assert syntactic_eq_alpha(actual, term) + @pytest.mark.parametrize("monoid", ALL_MONOIDS) @given(data=st.data()) @@ -95,7 +232,7 @@ def test_associativity(monoid, backend: Backend, data): with handler(NormalizeIntp): left = monoid.plus(monoid.plus(a, b), c) right = monoid.plus(a, monoid.plus(b, c)) - assert backend.eq(left, right) + assert syntactic_eq(left, right) @pytest.mark.parametrize("monoid", ALL_MONOIDS) @@ -123,7 +260,7 @@ def test_commutativity(monoid, backend: Backend, data): a = data.draw(backend.strategy(ret="scalar"))() b = data.draw(backend.strategy(ret="scalar"))() with handler(NormalizeIntp): - assert backend.eq(monoid.plus(a, b), monoid.plus(b, a)) + assert syntactic_eq(monoid.plus(a, b), monoid.plus(b, a)) @pytest.mark.parametrize("monoid", IDEMPOTENT) @@ -224,58 +361,83 @@ def test_plus_mapping(monoid, backend: Backend): backend.check_rewrite(lhs=lhs, rhs=rhs, rule=MonoidOverMapping()) -def test_plus_distributes(backend: Backend): +def test_plus_distributes_1(backend: Backend): + a, b, c = backend.define_vars("a", "b", "c", ret="scalar") + lhs = Product.plus(c(), Sum.plus(a(), b())) + rhs = Sum.plus(Product.plus(c(), a()), Product.plus(c(), b())) + backend.check_rewrite( + lhs=lhs, + rhs=rhs, + rule=coproduct(PlusDistr(), coproduct(PlusSingle(), PlusAssoc())), + ) + + +def test_plus_distributes_2(backend: Backend): + a, b, c = backend.define_vars("a", "b", "c", ret="scalar") + lhs = Product.plus(Sum.plus(a(), b()), c()) + rhs = Sum.plus(Product.plus(a(), c()), Product.plus(b(), c())) + backend.check_rewrite( + lhs=lhs, + rhs=rhs, + rule=coproduct(PlusDistr(), coproduct(PlusSingle(), PlusAssoc())), + ) + + +def test_plus_distributes_3(backend: Backend): a, b, c, d = backend.define_vars("a", "b", "c", "d", ret="scalar") lhs = Product.plus(Sum.plus(a(), b()), Sum.plus(c(), d())) - rhs = Product.plus( - Sum.plus( - Product.plus(a(), c()), - Product.plus(a(), d()), - Product.plus(b(), c()), - Product.plus(b(), d()), - ) + rhs = Sum.plus( + Product.plus(a(), c()), + Product.plus(a(), d()), + Product.plus(b(), c()), + Product.plus(b(), d()), + ) + backend.check_rewrite( + lhs=lhs, + rhs=rhs, + rule=coproduct(PlusDistr(), coproduct(PlusSingle(), PlusAssoc())), + ) + + +def test_plus_distributes_4(backend: Backend): + a, b, c, d = backend.define_vars("a", "b", "c", "d", ret="scalar") + lhs = Product.plus(Sum.plus(a(), b()), c(), d()) + rhs = Sum.plus(Product.plus(a(), c(), d()), Product.plus(b(), c(), d())) + backend.check_rewrite( + lhs=lhs, + rhs=rhs, + rule=coproduct(PlusDistr(), coproduct(PlusSingle(), PlusAssoc())), ) - backend.check_rewrite(lhs=lhs, rhs=rhs, rule=PlusDistr()) def test_plus_distributes_constant(backend: Backend): a, b, c, d, e = backend.define_vars("a", "b", "c", "d", "e", ret="scalar") lhs = Product.plus(Sum.plus(a(), b()), Sum.plus(c(), d()), e()) - rhs = Product.plus( - e(), - Sum.plus( - Product.plus(a(), c()), - Product.plus(a(), d()), - Product.plus(b(), c()), - Product.plus(b(), d()), - ), + rhs = Sum.plus( + Product.plus(a(), c(), e()), + Product.plus(a(), d(), e()), + Product.plus(b(), c(), e()), + Product.plus(b(), d(), e()), + ) + backend.check_rewrite( + lhs=lhs, + rhs=rhs, + rule=coproduct(PlusDistr(), coproduct(PlusSingle(), PlusAssoc())), ) - backend.check_rewrite(lhs=lhs, rhs=rhs, rule=PlusDistr()) def test_plus_distributes_multiple(backend: Backend): a, b, c, d = backend.define_vars("a", "b", "c", "d", ret="scalar") - lhs = Sum.plus( - Min.plus(a(), b()), - Min.plus(c(), d()), - Max.plus(a(), b()), - Max.plus(c(), d()), + lhs = Sum.plus(Min.plus(a(), b()), Max.plus(c(), d())) + rhs = Min.plus( + Max.plus(Sum.plus(a(), c()), Sum.plus(a(), d())), + Max.plus(Sum.plus(b(), c()), Sum.plus(b(), d())), ) - rhs = Sum.plus( - Min.plus( - Sum.plus(a(), c()), - Sum.plus(a(), d()), - Sum.plus(b(), c()), - Sum.plus(b(), d()), - ), - Max.plus( - Sum.plus(a(), c()), - Sum.plus(a(), d()), - Sum.plus(b(), c()), - Sum.plus(b(), d()), - ), + backend.check_rewrite( + lhs=lhs, + rhs=rhs, + rule=coproduct(PlusDistr(), coproduct(PlusSingle(), PlusAssoc())), ) - backend.check_rewrite(lhs=lhs, rhs=rhs, rule=PlusDistr()) @pytest.mark.parametrize("monoid", IDEMPOTENT) @@ -290,21 +452,20 @@ def test_plus_idempotent_consecutive(monoid, backend: Backend): @pytest.mark.parametrize("monoid", IDEMPOTENT) def test_plus_idempotent_non_consecutive(monoid, backend: Backend): - """``a, b, a`` — Semilattice (Min/Max) collapses via commutative - PlusDups.""" a, b = backend.define_vars("a", "b", ret="scalar") lhs = monoid.plus(a(), b(), a()) - rhs = monoid.plus(a(), b()) - backend.check_rewrite(lhs=lhs, rhs=rhs, rule=PlusDups()) + rhs = monoid.plus(a(), b(), a()) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=PlusConsecutiveDups()) @pytest.mark.parametrize("monoid", [Min, Max]) def test_plus_commutative_idempotent_long(monoid, backend: Backend): """Long alternation collapses via commutative dedup (Min/Max only).""" - a, b = backend.define_vars("a", "b", ret="scalar") - lhs = monoid.plus(a(), b(), a(), b(), b(), a(), a()) - rhs = monoid.plus(a(), b()) - backend.check_rewrite(lhs=lhs, rhs=rhs, rule=PlusDups()) + lhs = monoid.plus(0, 1, 0, 1, 1, 0, 0) + rhs = monoid.plus(0, 1) + backend.check_rewrite( + lhs=lhs, rhs=rhs, rule=coproduct(PlusOrder(), PlusConsecutiveDups()) + ) @pytest.mark.parametrize("monoid", WITH_ZERO) @@ -444,18 +605,138 @@ def test_reduce_reduce(monoid, backend: Backend): @pytest.mark.parametrize("monoid", COMMUTATIVE) -def test_reduce_plus(monoid, backend: Backend): - a, b = backend.define_vars("a", "b", ret="scalar") - A, B = backend.define_vars("A", "B", ret="stream") +def test_reduce_split_subset(monoid, backend: Backend): + """ReduceSplit confines a stream to the summands that use it: ``a`` is used + only by the first summand, so it is pushed into a reduce over just that + summand (the constant summand keeps its -- now innermost -- ``a`` reduce). + """ + a = backend.define_vars("a", ret="scalar") + A = backend.define_vars("A", ret="stream") + c = backend.define_vars("c", ret="scalar") + f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="scalar") - lhs = monoid.reduce(monoid.plus(a(), b()), {a: A(), b: B()}) + lhs = monoid.reduce(monoid.plus(f(a()), c()), {a: A()}) rhs = monoid.plus( - monoid.reduce(a(), {a: A(), b: B()}), - monoid.reduce(b(), {a: A(), b: B()}), + monoid.reduce(f(a()), {a: A()}), + monoid.reduce(c(), {a: A()}), ) backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceSplit()) +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_reduce_mask_hoist(monoid): + """A reduce-stream-independent mask condition lifts out of the reduce: + ``M.reduce(M.mask(v, c), {a: A}) == M.mask(M.reduce(v, {a: A}), c)``. + + Pinned to ``IntBackend``: ``Monoid.mask`` evaluates symbolically (via + ``MaskConcrete``) under the ops stack used by ``check_rewrite``; jax-array + masks are exercised separately in ``test_handlers_jax_monoid``. + """ + backend = IntBackend() + a, c, d = backend.define_vars("a", "c", "d", ret="scalar") + A = backend.define_vars("A", ret="stream") + f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="scalar") + + cond = c() == d() # independent of the reduced stream `a` + lhs = monoid.reduce(monoid.mask(f(a()), cond), {a: A()}) + rhs = monoid.mask(monoid.reduce(f(a()), {a: A()}), cond) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceMaskHoist()) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_reduce_mask_hoist_dependent_noop(monoid): + """The mask does NOT hoist when its condition depends on the reduced + stream (that case is gather territory for ``ReduceEqualityMaskRange``). + """ + backend = IntBackend() + a, c = backend.define_vars("a", "c", ret="scalar") + A = backend.define_vars("A", ret="stream") + f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="scalar") + + term = monoid.reduce(monoid.mask(f(a()), a() == c()), {a: A()}) + backend.check_rewrite(lhs=term, rhs=term, rule=ReduceMaskHoist()) + + +@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. + """ + 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(),)}, + ) + 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. + """ + 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") + + lhs = monoid.reduce( + monoid.mask(f(a()), And.plus(d() < e(), a() == c(), c() < e())), + {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(),)}, + ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceEqualityMaskRange()) + + +@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. + """ + 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()) + + +@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. + """ + 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) + ) + 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.reduce(monoid.mask(g(a()), c() == 0), {a: range(3)}), + ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceEqualityMaskRange()) + + 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") @@ -464,7 +745,7 @@ def test_reduce_independent_1(backend: Backend): rhs = Product.plus( Sum.reduce(Product.plus(a()), {a: A()}), Sum.reduce(Product.plus(b()), {b: B()}) ) - backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceFactorization()) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=Factor()) def test_reduce_independent_2(backend: Backend): @@ -482,7 +763,7 @@ def test_reduce_independent_2(backend: Backend): {b: B()}, ), ) - backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceFactorization()) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=Factor()) def test_reduce_independent_3_negative(backend: Backend): @@ -495,7 +776,7 @@ def test_reduce_independent_3_negative(backend: Backend): ) g = backend.define_vars("g", arg_types=(backend.scalar_typ,), ret="stream") - with handler(ReduceFactorization()): # ty:ignore[invalid-argument-type] + with handler(Factor()): # ty:ignore[invalid-argument-type] lhs = Sum.reduce( Product.plus(a(), b(), f(b(), c())), {a: A(), b: g(a()), c: C()} ) @@ -523,7 +804,7 @@ def test_reduce_independent_4(backend: Backend): {b: B()}, ), ) - backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceFactorization()) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=Factor()) def test_reduce_chain(backend: Backend): @@ -539,7 +820,7 @@ def test_reduce_chain(backend: Backend): Product.plus(h(y()), Sum.reduce(Product.plus(f(x()), g(x(), y())), {x: X()})), {y: Y()}, ) - backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceFactorization()) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=Factor()) @pytest.mark.parametrize("outer,inner", MONOID_PAIRS) @@ -562,7 +843,7 @@ def test_reduce_lift_shared(outer, inner, backend: Backend): ), {c: C()}, ) - backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceFactorization()) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=Factor()) @pytest.mark.parametrize("outer,inner", MONOID_PAIRS) @@ -590,80 +871,241 @@ def test_reduce_lift_shared_deps(outer, inner, backend: Backend): ), {c: C(), d: h(c())}, ) - backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceFactorization()) - - -def test_reduce_cartesian_3(): - backend = JaxBackend() - i = backend.define_vars("i", ret="scalar") - - with handler(NormalizeIntp): - value = CartesianProduct.reduce(jnp.zeros(2), {i: jnp.arange(3)}) - assert value.shape == (2**3, 3) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=Factor()) - with handler(NormalizeIntp): - value = CartesianProduct.reduce(jnp.zeros(2), {i: jnp.arange(1)}) - assert value.shape == (2**1, 1) - with handler(NormalizeIntp): - value = CartesianProduct.reduce(jnp.zeros(1), {i: jnp.arange(3)}) - assert value.shape == (1**3, 3) +def test_cartesian_union(): + i, a = Operation.define(int), Operation.define(int) + lhs = CartesianProduct.reduce( + Union.reduce([as_dict(((i(),), a()))], {a: range(2)}), {i: range(2)} + ) + with handler(NormalizeIntp), handler(EvaluateIntp): + rhs = evaluate(lhs) + assert syntactic_eq( + rhs, + [ + {(0,): 0, (1,): 0}, + {(0,): 0, (1,): 1}, + {(0,): 1, (1,): 0}, + {(0,): 1, (1,): 1}, + ], + ) @pytest.mark.parametrize("outer,inner", MONOID_PAIRS) def test_reduce_lifted_1(outer, inner, backend: Backend): a, i = backend.define_vars("a", "i", ret="scalar") - A, N, A_domain = backend.define_vars("A", "N", "A_domain", ret="stream") + N, A_domain = backend.define_vars("N", "A_domain", ret="stream") f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="scalar") + A = Operation.define(Mapping[tuple, backend.scalar_typ]) # type: ignore[name-defined] lhs = outer.reduce( - inner.reduce(f(a()), {a: A()}), - {A: CartesianProduct.reduce(A_domain(), {i: N()})}, + inner.reduce(f(A()[(i(),)]), {i: range(3)}), + { + A: CartesianProduct.reduce( + Union.reduce([as_dict(((i(),), a()))], {a: A_domain()}), {i: range(3)} + ) + }, + ) + rhs = inner.reduce(outer.reduce(f(a()), {a: A_domain()}), {i: range(3)}) + backend.check_rewrite( + lhs=lhs, + rhs=rhs, + rule=coproduct( + ReduceDistributeCartesianProduct(), + coproduct(ReduceUnion(), EliminateSingletonStreams()), + ), ) - rhs = inner.reduce(outer.reduce(inner.plus(f(a())), {a: A_domain()}), {i: N()}) - backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceDistributeCartesianProduct()) -def test_reduce_cartesian_1(): - backend = IntBackend() - a, i = backend.define_vars("a", "i", ret="scalar") - A = backend.define_vars("A", ret="stream") +def test_reduce_lifted_aligns_equal_ranges_by_row_position(): + """Equal-sized axes are aligned by row position, not stream order.""" + backend = JaxBackend() + a, p, q, i, j, k, l = backend.define_vars( + "a", "p", "q", "i", "j", "k", "l", ret="scalar" + ) + A = Operation.define(Mapping[tuple, backend.scalar_typ], name="A") + f, g = backend.define_vars("f", "g", arg_types=(backend.scalar_typ,), ret="scalar") + plate = range(2) - with handler(NormalizeIntp): - term1 = Sum.reduce( - Product.reduce(a(), {a: []}), - {A: CartesianProduct.reduce([], {i: []})}, - ) - term2 = Product.reduce(Sum.reduce(a(), {a: []}), {i: []}) - assert term1 == term2 + lhs = Sum.reduce( + Product.plus( + Product.reduce(f(A()[(i(), j())]), {j: plate, i: plate}), + Product.reduce(g(A()[(k(), l())]), {k: plate, l: plate}), + ), + { + A: CartesianProduct.reduce( + Union.reduce([as_dict(((p(), q()), a()))], {a: range(3)}), + {p: plate, q: plate}, + ) + }, + ) + rhs = Product.reduce( + Product.reduce( + Sum.reduce(Product.plus(f(a()), g(a())), {a: range(3)}), {q: plate} + ), + {p: plate}, + ) + norm = handler(ReduceDistributeCartesianProduct())(evaluate)(lhs) + assert syntactic_eq_alpha(norm, rhs) -def test_reduce_cartesian_2(): - backend = IntBackend() - a, i = backend.define_vars("a", "i", ret="scalar") - A = backend.define_vars("A", ret="stream") - with handler(NormalizeIntp): - term1 = Sum.reduce( - Product.reduce(a(), {a: A()}), - {A: CartesianProduct.reduce([(0,)], {i: [0]})}, - ) - term2 = Product.reduce(Sum.reduce(a(), {a: [0]}), {i: [0]}) - assert term1 == term2 +def test_reduce_lifted_unfactors_single_product_factor(): + """A unary product body can be exposed without an explicit Product.plus.""" + backend = JaxBackend() + a, b, i, p, h = backend.define_vars("a", "b", "i", "p", "h", ret="scalar") + A = Operation.define(Mapping[tuple, backend.scalar_typ], name="A") + f = backend.define_vars( + "f", + arg_types=(backend.scalar_typ, backend.scalar_typ), + ret="scalar", + ) + plate = range(2) + + lhs = Sum.reduce( + Product.plus( + Sum.reduce(Product.reduce(f(A()[(i(),)], b()), {i: plate}), {b: range(4)}), + h(), + ), + { + A: CartesianProduct.reduce( + Union.reduce([as_dict(((p(),), a()))], {a: range(3)}), {p: plate} + ) + }, + ) + rhs = Sum.reduce( + Product.plus( + Product.reduce(Sum.reduce(f(a(), b()), {a: range(3)}), {p: plate}), h() + ), + {b: range(4)}, + ) + + norm = handler(ReduceDistributeCartesianProduct())(evaluate)(lhs) + assert syntactic_eq_alpha(norm, rhs) + + +def test_reduce_lifted_sum_of_products_cartesian_body(): + """A cartesian-product reduction remains visible after SOP expansion.""" + backend = JaxBackend() + a, b, c_v, d_v, e_v, i, j = backend.define_vars( + "a", "b", "c_v", "d_v", "e_v", "i", "j", ret="scalar" + ) + d = Operation.define(Mapping[tuple, backend.scalar_typ], name="d") + f, f0, f1, f2 = backend.define_vars("f", "f0", "f1", "f2", ret="scalar") + + lhs = Sum.reduce( + Product.plus( + Product.reduce( + Sum.reduce( + f2()[(d()[(i(),)], e_v(), f()[(i(), j())], i(), j())], + {e_v: range(6)}, + ), + {i: range(2), j: range(3)}, + ), + Sum.reduce( + Product.plus( + Product.reduce( + Sum.reduce( + f1()[(b(), c_v(), d()[(i(),)], i())], + {c_v: range(4)}, + ), + {i: range(2)}, + ), + Sum.reduce(f0()[(a(), b())], {a: range(2)}), + ), + {b: range(3)}, + ), + ), + { + d: CartesianProduct.reduce( + Union.reduce([as_dict(((i(),), d_v()))], {d_v: range(5)}), + {i: range(2)}, + ) + }, + ) + + norm = handler(ReduceDistributeCartesianProduct())(evaluate)(lhs) + assert CartesianProduct.reduce not in fvsof(norm) + + product_streams = [] + + def walk(expr): + if isinstance(expr, Term): + if expr.op is Product.reduce: + product_streams.append(expr.args[1]) + for arg in expr.args: + walk(arg) + for arg in expr.kwargs.values(): + walk(arg) + elif isinstance(expr, tuple | list): + for item in expr: + walk(item) + elif isinstance(expr, Mapping): + for item in expr.values(): + walk(item) + + walk(norm) + assert [stream for streams in product_streams for stream in streams.values()] == [ + range(2), + range(3), + ] + + +def test_reduce_lifted_3(backend: Backend): + a, b, i, j = backend.define_vars("a", "b", "i", "j", ret="scalar") + N, A_domain, B_domain = backend.define_vars( + "N", "A_domain", "B_domain", ret="stream" + ) + f = backend.define_vars( + "f", arg_types=(backend.scalar_typ, backend.scalar_typ), ret="scalar" + ) + A = Operation.define(Mapping[tuple, backend.scalar_typ], name="A") # type: ignore[name-defined] + B = Operation.define(Mapping[tuple, backend.scalar_typ], name="B") # type: ignore[name-defined] + + lhs = Sum.reduce( + Product.reduce(f(A()[(i(), j())], B()[(i(), j())]), {i: range(2), j: range(3)}), + { + A: CartesianProduct.reduce( + Union.reduce([as_dict(((i(), j()), a()))], {a: A_domain()}), + {i: range(2), j: range(3)}, + ), + B: CartesianProduct.reduce( + Union.reduce([as_dict(((i(), j()), b()))], {b: B_domain()}), + {i: range(2), j: range(3)}, + ), + }, + ) + rhs = Product.reduce( + Product.reduce( + Sum.reduce(Sum.reduce(f(a(), b()), {a: A_domain()}), {b: B_domain()}), + {j: range(3)}, + ), + {i: range(2)}, + ) + norm = handler(ReduceDistributeCartesianProduct())(evaluate)(lhs) + assert syntactic_eq_alpha(norm, rhs) @pytest.mark.parametrize("outer,inner", MONOID_PAIRS) def test_reduce_lifted_multi_index(outer, inner, backend: Backend): a, i, j = backend.define_vars("a", "i", "j", ret="scalar") - A, N, M, A_domain = backend.define_vars("A", "N", "M", "A_domain", ret="stream") + N, M, A_domain = backend.define_vars("N", "M", "A_domain", ret="stream") + A = Operation.define(Mapping[tuple, backend.scalar_typ]) # type: ignore[name-defined] f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="scalar") lhs = outer.reduce( - inner.reduce(f(a()), {a: A()}), - {A: CartesianProduct.reduce(A_domain(), {i: N(), j: M()})}, + inner.reduce(f(A()[(i(), j())]), {i: range(3), j: range(2)}), + { + A: CartesianProduct.reduce( + Union.reduce([as_dict(((i(), j()), a()))], {a: A_domain()}), + {i: range(3), j: range(2)}, + ) + }, ) rhs = inner.reduce( - outer.reduce(inner.plus(f(a())), {a: A_domain()}), {i: N(), j: M()} + inner.reduce(outer.reduce(f(a()), {a: A_domain()}), {j: range(2)}), + {i: range(3)}, ) backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceDistributeCartesianProduct()) @@ -675,7 +1117,8 @@ def test_reduce_lifted_2(outer, inner, backend: Backend): """ a, i, s, t = backend.define_vars("a", "i", "s", "t", ret="scalar") - A, N, T = backend.define_vars("A", "N", "T", ret="stream") + N, T = backend.define_vars("N", "T", ret="stream") + A = Operation.define(Mapping[tuple, backend.scalar_typ]) # type: ignore[name-defined] A_domain = backend.define_vars( "A_domain", arg_types=(backend.scalar_typ,), ret="stream" ) @@ -684,19 +1127,32 @@ def test_reduce_lifted_2(outer, inner, backend: Backend): ) lhs = outer.reduce( - inner.reduce(inner.plus(f1(a(), s()), f2(t(), a())), {a: A()}), - {A: CartesianProduct.reduce(A_domain(i()), {i: N()}), t: T()}, + inner.reduce( + inner.plus(f1(A()[(i(),)], s()), f2(t(), A()[(i(),)])), {i: range(3)} + ), + { + A: CartesianProduct.reduce( + Union.reduce([as_dict(((i(),), a()))], {a: A_domain(i())}), + {i: range(3)}, + ), + t: T(), + }, ) rhs = outer.reduce( inner.reduce( - outer.reduce( - inner.plus(inner.plus(f1(a(), s()), f2(t(), a()))), {a: A_domain(i())} - ), - {i: N()}, + outer.reduce(inner.plus(f1(a(), s()), f2(t(), a())), {a: A_domain(i())}), + {i: range(3)}, ), {t: T()}, ) - backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceDistributeCartesianProduct()) + backend.check_rewrite( + lhs=lhs, + rhs=rhs, + rule=coproduct( + ReduceDistributeCartesianProduct(), + coproduct(ReduceUnion(), EliminateSingletonStreams()), + ), + ) # --------------------------------------------------------------------------- @@ -725,7 +1181,7 @@ def test_reduce_weighted_factorization(backend: Backend): Sum.reduce(f(a)*g(b), {a: Product.weighted(A, a, w_a), b: Product.weighted(B, b, w_b)}) = (Sum.reduce(w_a(a)*f(a), {a: A})) * (Sum.reduce(w_b(b)*g(b), {b: B})) - Exercises chaining of ``ReduceWeightedStream`` with ``ReduceFactorization`` + Exercises chaining of ``ReduceWeightedStream`` with ``Factor`` inside ``NormalizeIntp``. """ a, b = backend.define_vars("a", "b", ret="scalar") @@ -743,69 +1199,7 @@ def test_reduce_weighted_factorization(backend: Backend): Sum.reduce(Product.plus(w_b(b()), Product.plus(g(b()))), {b: B()}), ) backend.check_rewrite( - lhs=lhs, rhs=rhs, rule=coproduct(ReduceWeightedStream(), ReduceFactorization()) - ) - - -def test_reduce_cartesian_weighted_stream(backend: Backend): - """``CartesianProduct.reduce`` over a ``WeightedStream`` body whose weight - is independent of the plate var rewrites to a single joint - ``WeightedStream``: - - CartesianProduct.reduce(M.weighted(s, e, w(e)), {p: P}) - = M.weighted(CartesianProduct.reduce(s, {p: P}), row, M.reduce(w(e), {e: row()})) - """ - p, e_var = backend.define_vars("p", "e_var", ret="scalar") - S, P = backend.define_vars("S", "P", ret="stream") - w = backend.define_vars("w", arg_types=(backend.scalar_typ,), ret="scalar") - - lhs = CartesianProduct.reduce(Product.weighted(S(), w), {p: P()}) - row_var = Operation.define(Iterable[backend.scalar_typ], name="row") # type: ignore[name-defined] - rhs = Product.weighted( - CartesianProduct.reduce(S(), {p: P()}), - deffn(Product.reduce(w(e_var()), {e_var: row_var()}), row_var), - ) - backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceCartesianWeightedStream()) - - -def test_lift_weighted_cartesian(backend: Backend): - """Compose ``ReduceCartesianWeightedStream`` + ``ReduceWeightedStream`` + - ``ReduceDistributeCartesianProduct`` on a Sum-of-Product-of-weighted shape: - - Sum.reduce( - Product.reduce(body(a()), {a: A()}), - {A: CartesianProduct.reduce(Product.weighted(S, e, w(e)), {p: P})}, - ) - - The inner ``weighted`` becomes a joint ``weighted`` (rule 1), lifts its - per-element weight into the outer Sum body (rule 2), and the lifted form - matches the inversion pattern (rule 3), yielding:: - - Product.reduce( - Sum.reduce(Product.plus(w(a()), body(a())), {a: S}), - {p: P}, - ) - """ - a, p = backend.define_vars("a", "p", ret="scalar") - A, S, P = backend.define_vars("A", "S", "P", ret="stream") - body, w = backend.define_vars( - "body", "w", arg_types=(backend.scalar_typ,), ret="scalar" - ) - - lhs = Sum.reduce( - Product.reduce(body(a()), {a: A()}), - {A: CartesianProduct.reduce(Product.weighted(S(), w), {p: P()})}, - ) - rhs = Product.reduce( - Sum.reduce(Product.plus(w(a()), body(a())), {a: S()}), {p: P()} - ) - backend.check_rewrite( - lhs=lhs, - rhs=rhs, - rule=coproduct( - coproduct(ReduceWeightedStream(), ReduceCartesianWeightedStream()), - ReduceDistributeCartesianProduct(), - ), + lhs=lhs, rhs=rhs, rule=coproduct(ReduceWeightedStream(), Factor()) ) @@ -832,7 +1226,203 @@ def _f(v: int) -> float: w = Operation.define(_w, name="w") f = Operation.define(_f, name="f") - with handler(NormalizeIntp): + with handler(NormalizeIntp), handler(EvaluateIntp): result = evaluate(Sum.reduce(f(a()), {a: Product.weighted([1, 2, 3, 4], w)})) assert math.isclose(result, 10.0) + + +# --------------------------------------------------------------------------- +# CartesianProduct.plus (pure-Python ``CartesianProductPlus`` implementation) +# --------------------------------------------------------------------------- +# +# A ``CartesianProduct`` value is a list of "rows", each row a ``dict`` mapping +# index variables to values. ``plus`` takes the cartesian product of its +# argument lists, disjoint-merging the dicts of each combination into a single +# row. The identity is ``[{}]`` (one empty row) and the zero is ``[]`` (no +# rows). + + +@pytest.fixture +def cprod(): + """A handler scope in which ``CartesianProduct.plus`` is concrete.""" + with handler(CartesianProductPlus()): + yield + + +def test_cprod_plus_two_singletons(cprod): + """Two single-row lists merge into one row with the union of their keys.""" + assert CartesianProduct.plus([{"a": 1}], [{"b": 2}]) == [{"a": 1, "b": 2}] + + +def test_cprod_plus_multi_key_rows(cprod): + """Every key of a multi-key row survives the merge (regression: the merge + used to keep only the last key of each dict).""" + assert CartesianProduct.plus([{"a": 1, "b": 2}], [{"c": 3, "d": 4}]) == [ + {"a": 1, "b": 2, "c": 3, "d": 4} + ] + + +def test_cprod_plus_cartesian_expansion(cprod): + """The result enumerates the full cartesian product of the input rows.""" + result = CartesianProduct.plus([{"a": 1}, {"a": 2}, {"a": 3}], [{"b": 4}, {"b": 5}]) + assert result == [ + {"a": 1, "b": 4}, + {"a": 1, "b": 5}, + {"a": 2, "b": 4}, + {"a": 2, "b": 5}, + {"a": 3, "b": 4}, + {"a": 3, "b": 5}, + ] + + +def test_cprod_plus_cardinality(cprod): + """|plus(A, B, C)| == |A| * |B| * |C|.""" + a = [{"a": i} for i in range(2)] + b = [{"b": i} for i in range(3)] + c = [{"c": i} for i in range(4)] + assert len(CartesianProduct.plus(a, b, c)) == 2 * 3 * 4 + + +def test_cprod_plus_three_args(cprod): + """``plus`` is variadic and merges across all arguments at once.""" + assert CartesianProduct.plus([{"a": 1}], [{"b": 2}], [{"c": 3}]) == [ + {"a": 1, "b": 2, "c": 3} + ] + + +def test_cprod_plus_single_arg(cprod): + """A single argument is returned row-for-row (product of one factor).""" + assert CartesianProduct.plus([{"a": 1}, {"a": 2}]) == [{"a": 1}, {"a": 2}] + + +def test_cprod_plus_identity_right(cprod): + """The identity ``[{}]`` is a right unit: merging an empty row changes + nothing.""" + assert CartesianProduct.plus([{"a": 1}, {"a": 2}], CartesianProduct.identity) == [ + {"a": 1}, + {"a": 2}, + ] + + +def test_cprod_plus_identity_left(cprod): + """The identity ``[{}]`` is a left unit.""" + assert CartesianProduct.plus(CartesianProduct.identity, [{"a": 1}, {"a": 2}]) == [ + {"a": 1}, + {"a": 2}, + ] + + +def test_cprod_plus_identity_with_identity(cprod): + """Identity ⊕ identity == identity.""" + assert ( + CartesianProduct.plus(CartesianProduct.identity, CartesianProduct.identity) + == CartesianProduct.identity + ) + + +def test_cprod_plus_zero_right(cprod): + """The zero ``[]`` absorbs on the right (empty cartesian product).""" + assert CartesianProduct.plus([{"a": 1}], CartesianProduct.zero) == [] + + +def test_cprod_plus_zero_left(cprod): + """The zero ``[]`` absorbs on the left.""" + assert CartesianProduct.plus(CartesianProduct.zero, [{"a": 1}]) == [] + + +def test_cprod_plus_zero_among_many(cprod): + """A single zero factor anywhere collapses the whole product to ``[]``.""" + assert CartesianProduct.plus([{"a": 1}], [], [{"c": 3}]) == [] + + +def test_cprod_plus_empty_rows_preserved(cprod): + """Rows that are themselves empty dicts merge cleanly.""" + assert CartesianProduct.plus([{}], [{"a": 1}]) == [{"a": 1}] + assert CartesianProduct.plus([{}], [{}]) == [{}] + + +def test_cprod_plus_associative(cprod): + """plus(plus(A, B), C) == plus(A, plus(B, C)) == plus(A, B, C).""" + a = [{"a": 1}, {"a": 2}] + b = [{"b": 3}] + c = [{"c": 4}, {"c": 5}] + flat = CartesianProduct.plus(a, b, c) + assert CartesianProduct.plus(CartesianProduct.plus(a, b), c) == flat + assert CartesianProduct.plus(a, CartesianProduct.plus(b, c)) == flat + + +def test_cprod_plus_duplicate_key_raises(cprod): + """Merging rows that share a key is ill-defined and rejected.""" + with pytest.raises(ValueError, match="Duplicate key found: 'a'"): + CartesianProduct.plus([{"a": 1}], [{"a": 2}]) + + +def test_cprod_plus_duplicate_key_in_multi_key_row_raises(cprod): + """The duplicate-key check sees every key of a multi-key row, not just the + last (regression for the same merge bug).""" + with pytest.raises(ValueError, match="Duplicate key found: 'a'"): + CartesianProduct.plus([{"a": 1, "x": 9}], [{"a": 2}]) + + +def test_cprod_plus_does_not_mutate_inputs(cprod): + """The input rows are left untouched; merges build fresh dicts.""" + left = [{"a": 1}] + right = [{"b": 2}] + CartesianProduct.plus(left, right) + assert left == [{"a": 1}] + assert right == [{"b": 2}] + + +def test_cprod_plus_forwards_on_term(): + """A symbolic (``Term``) argument cannot be enumerated, so ``plus`` forwards + and the call builds an unevaluated ``Term`` instead of a value.""" + x = Operation.define(Iterable, name="x") + with handler(CartesianProductPlus()): + result = CartesianProduct.plus(x(), [{"a": 1}]) + assert isinstance(result, Term) + assert result.op is CartesianProduct.plus + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_reduce_dependent_range_mask(monoid, backend: Backend): + """A dependent range stream gets rewritten to the referent's bbox stream, + with the original constraint folded into the body as a where-guard. + + reduce(M, {u: range(0, N, 1), v: range(0, u(), 1)}, body) + ≡ reduce(M, {u: range(0, N, 1), v: range(0, N, 1)}, where(v() < u(), body, M.identity)) + """ + (u, v) = backend.define_vars("u", "v", ret="scalar") + N = 5 + f = backend.define_vars( + "f", arg_types=[backend.scalar_typ, backend.scalar_typ], ret="scalar" + ) + + body = f(u(), v()) + + lhs = monoid.reduce(body, {u: range_(N), v: range_(u())}) + rhs = monoid.reduce(monoid.mask(body, v() < u()), {u: range_(N), v: range_(N)}) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceDependentRangeMask()) + + +@pytest.mark.parametrize("Sum,Product", COMMUTATIVE_MONOID_PAIRS) +def test_reduce_unfactor_simple(Sum, Product, backend: Backend): + x, y, g = backend.define_vars("x", "y", "g", ret="scalar") + X, Y = backend.define_vars("X", "Y", ret="stream") + f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="scalar") + lhs = Sum.reduce(Product.plus(Sum.reduce(f(x()), {x: X()}), g()), {y: Y()}) + rhs = Sum.reduce(Product.plus(f(x()), g()), {x: X(), y: Y()}) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceUnfactor()) + + +@pytest.mark.parametrize("Sum,Product", COMMUTATIVE_MONOID_PAIRS) +def test_reduce_unfactor_reduces(Sum, Product, backend: Backend): + x, y, z = backend.define_vars("x", "y", "z", ret="scalar") + X, Y, Z = backend.define_vars("X", "Y", "Z", ret="stream") + f, g = backend.define_vars("f", "g", arg_types=(backend.scalar_typ,), ret="scalar") + lhs = Sum.reduce( + Product.plus(Sum.reduce(f(x()), {x: X()}), Sum.reduce(g(y()), {y: Y()})), + {z: Z()}, + ) + rhs = Sum.reduce(Product.plus(f(x()), g(y())), {x: X(), y: Y(), z: Z()}) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceUnfactor()) diff --git a/tests/test_ops_semantics.py b/tests/test_ops_semantics.py index 04e595f9e..66c391631 100644 --- a/tests/test_ops_semantics.py +++ b/tests/test_ops_semantics.py @@ -24,6 +24,7 @@ deffn, defop, implements, + syntactic_eq, ) from effectful.ops.types import Interpretation, NotHandled, Operation, Term @@ -469,7 +470,9 @@ def Nested(*args, **kwargs): t = Nested([{"a": y()}, x(), (x(), y())], x(), arg1={"b": x()}) with handler({x: lambda: 1, y: lambda: 2}): - assert evaluate(t) == Nested([{"a": 2}, 1, (1, 2)], 1, arg1={"b": 1}) + assert syntactic_eq( + evaluate(t), Nested([{"a": 2}, 1, (1, 2)], 1, arg1={"b": 1}) + ) def test_memoized_interpretation(): From 76896e3533f428202e5b394610f28e42ddb2985f Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Sat, 25 Jul 2026 14:51:34 -0400 Subject: [PATCH 12/16] Monoid.reduce(b, {}) == b (#723) --- effectful/handlers/jax/monoid.py | 10 ++- effectful/ops/monoid.py | 46 ++++++------- tests/test_ops_monoid.py | 113 +++++++------------------------ 3 files changed, 48 insertions(+), 121 deletions(-) diff --git a/effectful/handlers/jax/monoid.py b/effectful/handlers/jax/monoid.py index 7cf369894..3acf89f88 100644 --- a/effectful/handlers/jax/monoid.py +++ b/effectful/handlers/jax/monoid.py @@ -363,9 +363,7 @@ def _inequality_to_scan(cmp_op, args, tail_mask_elems): jax_getitem(scan_val, (index,)), And.plus(*tail_mask_elems) ) tail_streams = {k: v for (k, v) in streams.items() if k != stream_op} - if tail_streams: - return monoid.reduce(tail_body, tail_streams) - return tail_body + return monoid.reduce(tail_body, tail_streams) mask_elems = _conjuncts(mask) for i, elem in enumerate(mask_elems): @@ -508,8 +506,8 @@ def _(self, body, streams: Streams): # create leading reduction dimensions index = tuple(k() for k in streams) - pos_lhs = Sum.reduce(Sum.delta(index, lhs), streams) - pos_rhs = Sum.reduce(Sum.delta(index, rhs), streams) + pos_lhs = Sum.reduce(Sum.delta(index, lhs) if index else lhs, streams) + pos_rhs = Sum.reduce(Sum.delta(index, rhs) if index else rhs, streams) dims = "".join(get_symbol(i) for i in range(len(streams))) contraction = jnp.einsum(f"{dims}...,{dims}...->...", pos_lhs, pos_rhs) @@ -745,7 +743,7 @@ def term(self) -> Expr[Callable]: dims = [(self.out_vars[c], self.sizes[c]) for c in self.out_spec] reductions = self._build_plate_reductions(self.plate_tree) - reduction = Sum.reduce(reductions, streams) if streams else reductions + reduction = Sum.reduce(reductions, streams) return deffn( bind_dims( deffn(reduction, *(self.out_vars[c] for c in self.out_spec))( diff --git a/effectful/ops/monoid.py b/effectful/ops/monoid.py index 11d58ff55..4830133d3 100644 --- a/effectful/ops/monoid.py +++ b/effectful/ops/monoid.py @@ -143,8 +143,8 @@ def reduce[A, B, U: Body]( streams: Annotated[Streams, Scoped[A]], ) -> Annotated[U, Scoped[B]]: """Reduce ``body`` over ``streams``. Handlers supply per-monoid and - broadcasting behavior; the default rule only handles the empty-stream - case. + broadcasting behavior. + """ raise NotHandled @@ -561,11 +561,20 @@ def plus(self, monoid, *args): return monoid.plus(*(args[i] for i in dedup_args)) +class ReduceEmpty(ObjectInterpretation): + @implements(Monoid.reduce) + def _(self, monoid, body, streams): + if streams: + return fwd() + + return body + + class ReducePartial(ObjectInterpretation): @implements(Monoid.reduce) def _(self, monoid, body, streams): if not streams: - return monoid.identity + return fwd() for stream_key, stream_body, streams_tail in outer_stream(streams): if isinstance(stream_body, Term): @@ -984,8 +993,7 @@ def _getitem(mapping, idx1): inner_tail_streams = { k: v for (k, v) in inner_streams.items() if k != inner_plate_op } - if inner_tail_streams: - subst_body = inner_monoid.reduce(subst_body, inner_tail_streams) + subst_body = inner_monoid.reduce(subst_body, inner_tail_streams) combined_factors.append(subst_body) combined = ( @@ -1013,33 +1021,20 @@ def _to_body(mapping, key): peeled_body = Union.reduce( [as_dict((peeled_idx, union_body))], union_streams ) - if not peeled_cprod_streams: - peeled_cprod = peeled_body - else: - peeled_cprod = CartesianProduct.reduce( - peeled_body, peeled_cprod_streams - ) + peeled_cprod = CartesianProduct.reduce( + peeled_body, peeled_cprod_streams + ) inner_reduce_body = monoid.reduce(combined, {stream_key: peeled_cprod}) peeled_reduce = inner_monoid.reduce( - inner_reduce_body, - {shared_plate_op: plate_range}, - ) - - result_body = ( - inner_monoid.plus(peeled_reduce, *row_outer_factors) - if row_outer_factors - else peeled_reduce + inner_reduce_body, {shared_plate_op: plate_range} ) + result_body = inner_monoid.plus(peeled_reduce, *row_outer_factors) # Include any extra sum streams outermost. In particular, the # non-reduce product factors above remain outside the plate fold. tail_streams = {k: v for (k, v) in streams.items() if k != stream_key} - if tail_streams: - result = monoid.reduce(result_body, tail_streams) - else: - result = result_body - + result = monoid.reduce(result_body, tail_streams) return result return fwd() @@ -1352,7 +1347,7 @@ def reduce(self, monoid, body, streams): } # reduce over no streams is a single (empty) assignment, i.e. the body # itself -- not the monoid identity. - return monoid.reduce(new_body, new_streams) if new_streams else new_body + return monoid.reduce(new_body, new_streams) class WhereHoist(ObjectInterpretation): @@ -1752,6 +1747,7 @@ def extend(self, *intps: Interpretation) -> typing.Self: MonoidOverSequence(), MonoidOverMapping(), MonoidOverCallable(), + ReduceEmpty(), ReduceFusion(), ReduceUnion(), ReduceSplit(), diff --git a/tests/test_ops_monoid.py b/tests/test_ops_monoid.py index 41fc01aa4..a2d2df2e3 100644 --- a/tests/test_ops_monoid.py +++ b/tests/test_ops_monoid.py @@ -1,3 +1,4 @@ +import functools import math import sys import typing @@ -32,6 +33,7 @@ ReduceDependentRangeMask, ReduceDisjunctiveDisequalityMask, ReduceDistributeCartesianProduct, + ReduceEmpty, ReduceEqualityMaskRange, ReduceFusion, ReduceMaskHoist, @@ -539,7 +541,7 @@ def test_eliminate_singleton_only_stream(monoid, backend: Backend): f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="scalar") lhs = monoid.reduce(f(x()), {x: (a(),)}) - rhs = f(a()) + rhs = monoid.reduce(f(a()), {}) backend.check_rewrite(lhs=lhs, rhs=rhs, rule=EliminateSingletonStreams()) @@ -587,8 +589,8 @@ def test_reduce_no_streams(monoid, backend: Backend): a = backend.define_vars("a", ret="scalar") lhs = monoid.reduce(a(), {}) - rhs = monoid.identity - backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReducePartial()) + rhs = a() + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceEmpty()) @pytest.mark.parametrize("monoid", ALL_MONOIDS) @@ -892,6 +894,18 @@ def test_cartesian_union(): ) +_LiftingIntp = functools.reduce( + coproduct, # type: ignore + ( + ReduceDistributeCartesianProduct(), + ReduceUnion(), + EliminateSingletonStreams(), + ReduceEmpty(), + PlusSingle(), + ), +) + + @pytest.mark.parametrize("outer,inner", MONOID_PAIRS) def test_reduce_lifted_1(outer, inner, backend: Backend): a, i = backend.define_vars("a", "i", ret="scalar") @@ -908,14 +922,7 @@ def test_reduce_lifted_1(outer, inner, backend: Backend): }, ) rhs = inner.reduce(outer.reduce(f(a()), {a: A_domain()}), {i: range(3)}) - backend.check_rewrite( - lhs=lhs, - rhs=rhs, - rule=coproduct( - ReduceDistributeCartesianProduct(), - coproduct(ReduceUnion(), EliminateSingletonStreams()), - ), - ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=_LiftingIntp) def test_reduce_lifted_aligns_equal_ranges_by_row_position(): @@ -947,7 +954,7 @@ def test_reduce_lifted_aligns_equal_ranges_by_row_position(): {p: plate}, ) - norm = handler(ReduceDistributeCartesianProduct())(evaluate)(lhs) + norm = handler(_LiftingIntp)(evaluate)(lhs) assert syntactic_eq_alpha(norm, rhs) @@ -981,77 +988,10 @@ def test_reduce_lifted_unfactors_single_product_factor(): {b: range(4)}, ) - norm = handler(ReduceDistributeCartesianProduct())(evaluate)(lhs) + norm = handler(_LiftingIntp)(evaluate)(lhs) assert syntactic_eq_alpha(norm, rhs) -def test_reduce_lifted_sum_of_products_cartesian_body(): - """A cartesian-product reduction remains visible after SOP expansion.""" - backend = JaxBackend() - a, b, c_v, d_v, e_v, i, j = backend.define_vars( - "a", "b", "c_v", "d_v", "e_v", "i", "j", ret="scalar" - ) - d = Operation.define(Mapping[tuple, backend.scalar_typ], name="d") - f, f0, f1, f2 = backend.define_vars("f", "f0", "f1", "f2", ret="scalar") - - lhs = Sum.reduce( - Product.plus( - Product.reduce( - Sum.reduce( - f2()[(d()[(i(),)], e_v(), f()[(i(), j())], i(), j())], - {e_v: range(6)}, - ), - {i: range(2), j: range(3)}, - ), - Sum.reduce( - Product.plus( - Product.reduce( - Sum.reduce( - f1()[(b(), c_v(), d()[(i(),)], i())], - {c_v: range(4)}, - ), - {i: range(2)}, - ), - Sum.reduce(f0()[(a(), b())], {a: range(2)}), - ), - {b: range(3)}, - ), - ), - { - d: CartesianProduct.reduce( - Union.reduce([as_dict(((i(),), d_v()))], {d_v: range(5)}), - {i: range(2)}, - ) - }, - ) - - norm = handler(ReduceDistributeCartesianProduct())(evaluate)(lhs) - assert CartesianProduct.reduce not in fvsof(norm) - - product_streams = [] - - def walk(expr): - if isinstance(expr, Term): - if expr.op is Product.reduce: - product_streams.append(expr.args[1]) - for arg in expr.args: - walk(arg) - for arg in expr.kwargs.values(): - walk(arg) - elif isinstance(expr, tuple | list): - for item in expr: - walk(item) - elif isinstance(expr, Mapping): - for item in expr.values(): - walk(item) - - walk(norm) - assert [stream for streams in product_streams for stream in streams.values()] == [ - range(2), - range(3), - ] - - def test_reduce_lifted_3(backend: Backend): a, b, i, j = backend.define_vars("a", "b", "i", "j", ret="scalar") N, A_domain, B_domain = backend.define_vars( @@ -1083,7 +1023,7 @@ def test_reduce_lifted_3(backend: Backend): ), {i: range(2)}, ) - norm = handler(ReduceDistributeCartesianProduct())(evaluate)(lhs) + norm = handler(_LiftingIntp)(evaluate)(lhs) assert syntactic_eq_alpha(norm, rhs) @@ -1107,7 +1047,7 @@ def test_reduce_lifted_multi_index(outer, inner, backend: Backend): inner.reduce(outer.reduce(f(a()), {a: A_domain()}), {j: range(2)}), {i: range(3)}, ) - backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceDistributeCartesianProduct()) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=_LiftingIntp) @pytest.mark.parametrize("outer,inner", MONOID_PAIRS) @@ -1145,14 +1085,7 @@ def test_reduce_lifted_2(outer, inner, backend: Backend): ), {t: T()}, ) - backend.check_rewrite( - lhs=lhs, - rhs=rhs, - rule=coproduct( - ReduceDistributeCartesianProduct(), - coproduct(ReduceUnion(), EliminateSingletonStreams()), - ), - ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=_LiftingIntp) # --------------------------------------------------------------------------- From 533f8388c087c81f3540f0237c4a4b66a8753548 Mon Sep 17 00:00:00 2001 From: eb8680 Date: Tue, 28 Jul 2026 12:19:37 -0400 Subject: [PATCH 13/16] Don't route all-scalar monoid ops through jax (#729) `_jax_args` admitted `jax.typing.ArrayLike`, a union that includes `bool`, `int`, `float` and `complex`, so the jax `Monoid.plus` handlers claimed pure-Python scalar arithmetic. They extend `EvaluateIntp` after the scalar implementations and so take precedence, silently narrowing a Python float to a `float32` array and leaving downstream rules treating a scalar body as array-valued. Require at least one genuine array. Co-authored-by: Claude Opus 5 (1M context) --- effectful/handlers/jax/monoid.py | 19 ++++++++---- tests/test_handlers_jax_monoid.py | 49 +++++++++++++++++++++++++++++++ 2 files changed, 63 insertions(+), 5 deletions(-) diff --git a/effectful/handlers/jax/monoid.py b/effectful/handlers/jax/monoid.py index 3acf89f88..2e99d6859 100644 --- a/effectful/handlers/jax/monoid.py +++ b/effectful/handlers/jax/monoid.py @@ -68,12 +68,21 @@ def _jax_args(args): - """True iff ``args`` is non-empty and every arg is a concrete - :class:`jax.typing.ArrayLike` or named tensor. - + """True iff ``args`` is non-empty, every arg is a concrete + :class:`jax.typing.ArrayLike` or named tensor, and at least one of them is + an array rather than a Python scalar. + + :class:`jax.typing.ArrayLike` is a union that includes ``bool``, ``int``, + ``float`` and ``complex``, so admitting it alone would claim pure-Python + scalar arithmetic. These handlers extend ``EvaluateIntp`` after the scalar + implementations and so take precedence over them, which would silently + narrow a Python float to a ``float32`` array and leave downstream rules + treating a scalar body as array-valued. """ - return args and all( - isinstance(a, jax.typing.ArrayLike) or is_eager_array(a) for a in args + return ( + args + and all(isinstance(a, jax.typing.ArrayLike) or is_eager_array(a) for a in args) + and any(not isinstance(a, bool | int | float | complex) for a in args) ) diff --git a/tests/test_handlers_jax_monoid.py b/tests/test_handlers_jax_monoid.py index 14a1db08e..a538664d8 100644 --- a/tests/test_handlers_jax_monoid.py +++ b/tests/test_handlers_jax_monoid.py @@ -124,6 +124,55 @@ def test_reduce_array_2(monoid, reductor, backend: JaxBackend): assert jnp.allclose(actual, expected) +SCALAR_PLUS = [ + pytest.param(Sum, 3.0, id="Sum"), + pytest.param(Product, 2.0, id="Product"), +] + + +@pytest.mark.parametrize("monoid,expected", SCALAR_PLUS) +def test_plus_scalars_stays_scalar(monoid, expected): + """``Monoid.plus`` of plain Python numbers must not become a ``jax.Array``. + + ``jax.typing.ArrayLike`` is a union that includes ``bool``, ``int``, + ``float`` and ``complex``, so the jax ``plus`` handlers -- which extend + ``EvaluateIntp`` after the scalar implementations and therefore take + precedence -- must not claim pure-Python scalar arithmetic. + """ + with handler(NormalizeIntp), handler(EvaluateIntp): + actual = monoid.plus(1.0, 2.0) + + assert not isinstance(actual, jax.Array) + assert isinstance(actual, float) + assert actual == expected + + +@pytest.mark.parametrize( + "monoid,expected", + [pytest.param(Sum, 6.0, id="Sum"), pytest.param(Product, 8.0, id="Product")], +) +def test_reduce_scalar_body_stays_scalar(monoid, expected, backend: JaxBackend): + """A reduction over a scalar body is likewise plain Python arithmetic.""" + i = backend.define_vars("i", ret="scalar") + + with handler(NormalizeIntp), handler(EvaluateIntp): + actual = monoid.reduce(2.0, {i: range(3)}) + + assert not isinstance(actual, jax.Array) + assert isinstance(actual, float) + assert actual == expected + + +@pytest.mark.parametrize("monoid,expected", SCALAR_PLUS) +def test_plus_mixed_array_and_scalar_is_array(monoid, expected): + """One genuine array is enough: the narrowing must not be over-applied.""" + with handler(NormalizeIntp), handler(EvaluateIntp): + actual = monoid.plus(jnp.asarray([1.0, 1.0]), 2.0) + + assert isinstance(actual, jax.Array) + assert jnp.allclose(actual, jnp.asarray([expected, expected])) + + @pytest.mark.parametrize("monoid,reductor", MONOIDS) def test_arange_reduce_indirect(monoid, reductor, backend: JaxBackend): """When the range var is used both as a direct index and as a value From fe1e7606d0745ca97b746bb78c2d2d975c170005 Mon Sep 17 00:00:00 2001 From: eb8680 Date: Tue, 28 Jul 2026 15:54:32 -0400 Subject: [PATCH 14/16] Add a generator-expression bytecode disassembler (#725) * Don't route all-scalar monoid ops through jax `_jax_args` admitted `jax.typing.ArrayLike`, a union that includes `bool`, `int`, `float` and `complex`, so the jax `Monoid.plus` handlers claimed pure-Python scalar arithmetic. They extend `EvaluateIntp` after the scalar implementations and so take precedence, silently narrowing a Python float to a `float32` array and leaving downstream rules treating a scalar body as array-valued. Require at least one genuine array. Co-Authored-By: Claude Opus 5 (1M context) * Add a generator-expression bytecode disassembler `effectful/internals/disassembly.py` symbolically interprets the bytecode of a generator expression (and of lambdas and comprehensions nested inside it) back into an `ast` node, so a comprehension's source syntax can be recovered from the code object at runtime. Supports CPython 3.12 and 3.13. Standalone: imports nothing from `effectful` and touches no existing code. Co-Authored-By: Claude Opus 5 (1M context) * Address review comments on the generator-expression disassembler Six fixes, each with tests that fail without them: - `handle_build_map` read the key/value pairs of a dict display from the top of the stack down, reversing source order: a later duplicate key lost to an earlier one, and side effects ran backwards. - `_ensure_ast_tuple` treated any tuple whose first element was the string "dict_item" as an internal marker and dropped that element. Nothing produced such a marker; user data holding that string was silently corrupted. The special case is gone. - A free variable was reconstructed as a bare `ast.Name`, so evaluating the result resolved it against the evaluating namespace instead of the captured cell. The captured value is now written into the tree, for the generator itself, for lambdas reached as live objects, and for lambdas and comprehensions nested inside. A cell the comprehension creates -- a target captured by a nested lambda -- still stands as a name, since the reconstruction binds it too. A capture with no AST spelling, including an iterator, raises `TypeError` rather than reconstructing to a name that would answer differently. - `_ensure_ast_iterator_adaptor` ignored the strictness a `zip` pickles as reduction state, so a strict zip silently truncated ragged input where the original raised. - A lambda reached as a live object lost its default values, which live on the function rather than in its code object, leaving parameters with no way to be filled. - `disassemble` asserted on its input; it now raises `ValueError`, and checks the generator has not been started rather than leaving that to an assert further in. Also documents what reconstruction does and does not recover: evaluating the result re-runs every expression in it, so a stateful filter answers against state as it then stands. 663 passed, 2 xfailed on 3.12, 3.13 and 3.14. --------- Co-authored-by: Claude Opus 5 (1M context) --- effectful/internals/disassembly.py | 3590 ++++++++++++++++++++++++++ tests/test_internals_disassembler.py | 2109 +++++++++++++++ 2 files changed, 5699 insertions(+) create mode 100644 effectful/internals/disassembly.py create mode 100644 tests/test_internals_disassembler.py diff --git a/effectful/internals/disassembly.py b/effectful/internals/disassembly.py new file mode 100644 index 000000000..d2132b193 --- /dev/null +++ b/effectful/internals/disassembly.py @@ -0,0 +1,3590 @@ +""" +Generator expression bytecode reconstruction module. + +This module provides functionality to reconstruct AST representations from compiled +generator expressions by analyzing their bytecode. The primary use case is to recover +the original structure of generator comprehensions from their compiled form. + +The only public-facing interface is the `disassemble()` function, which takes a +generator object and returns an AST node representing the original comprehension. +All other functions and classes in this module are internal implementation details. + +Example: + >>> g = (x * 2 for x in range(10) if x % 2 == 0) + >>> ast_node = disassemble(g) + >>> # ast_node is now an ast.Expression representing the original expression + +What is recovered, and what is not: + + What comes back is the comprehension's *syntax*, not the generator's + suspended state, so evaluating the reconstruction runs every expression in + it a second time. For a comprehension over pure expressions the two agree + element for element. Where they can part company: + + * A **stateful filter or element expression** -- one that mutates something, + or whose value depends on state that has moved on -- is re-run, and can + answer differently the second time. `(x for x in xs if next(flags))` + reconstructs faithfully as source, but iterating the reconstruction draws + from `flags` where it now stands, not from where it stood when the + original generator was built. The side effects happen twice, once per + iteration of each generator. + + * The **outermost iterable** is not part of the comprehension's bytecode -- + it is an object the generator already holds -- so it is recovered from + that object rather than from source. What lands in the tree is a snapshot + of the elements not yet consumed: `iter([1, 2, 3])` advanced once becomes + the literal `[2, 3]`. The expression that produced it is gone, and so is + any laziness it had. + + * A **free variable** is likewise recovered by value, not by name: the value + in the closure cell is written into the tree. A captured object with no + AST spelling raises `TypeError` rather than reconstructing to a name that + would resolve against the evaluating namespace instead. +""" + +import ast +import builtins +import collections +import collections.abc +import copy +import dis +import enum +import functools +import inspect +import itertools +import sys +import types +import typing +from collections.abc import Callable, Generator, Iterator +from dataclasses import dataclass, field, replace + +CompExp = ast.GeneratorExp | ast.ListComp | ast.SetComp | ast.DictComp + + +class Placeholder(ast.Name): + """Placeholder for AST nodes that are not yet resolved.""" + + def __init__( + self, + id: typing.Literal[".PLACEHOLDER"] = ".PLACEHOLDER", + ctx: ast.Load = ast.Load(), + ): + super().__init__(id=id, ctx=ctx) + + +class DummyIterName(ast.Name): + """Dummy name for the iterator variable in generator expressions.""" + + def __init__(self, id: typing.Literal[".0"] = ".0", ctx: ast.Load = ast.Load()): + super().__init__(id=id, ctx=ctx) + + +class Skipped(ast.Name): + """Placeholder for skipped branches in if-expressions. + + ``id`` is defaulted so that ``copy.deepcopy`` can reconstruct the node: on + Python 3.12 ``ast.AST.__reduce__`` supplies no positional arguments (3.13+ + supplies one per field), so the constructor must be callable with none. + """ + + def __init__(self, id: str = "", ctx: ast.Load = ast.Load()): + super().__init__(id=id, ctx=ctx) + + +class CommonConstant(ast.Name): + """A constant pushed by 3.14's LOAD_COMMON_CONSTANT. + + 3.14 can inline `any()`/`all()` over a generator, guarding the fast path + with `loaded_name is `. Marking the builtin lets that guard be + recognised so the generic call path is followed instead, which still spells + out the original call. + """ + + def __init__(self, id: str = "", ctx: ast.Load = ast.Load()): + super().__init__(id=id, ctx=ctx) + + +class TargetHole(ast.Name): + """Placeholder for a comprehension loop target that is not yet named. + + ``FOR_ITER`` knows that a loop target exists but not what it is called; the + name only arrives with the ``STORE_FAST``/``UNPACK_SEQUENCE`` instructions + that follow. Each hole carries a unique ``id`` so that it can be located + again inside a comprehension after the surrounding state has been copied, + which matters for unpacking targets where several holes are live at once. + """ + + _counter: typing.ClassVar[Iterator[int]] = itertools.count() + + def __init__(self, id: str = "", ctx: ast.Store = ast.Store()): + super().__init__(id=id or f".TARGET_{next(TargetHole._counter)}", ctx=ctx) + + +class ReplaceTargetHole(ast.NodeTransformer): + """Replace the uniquely-identified :class:`TargetHole` ``id`` with ``replacement``.""" + + id: str + replacement: ast.expr + + def __init__(self, id: str, replacement: ast.expr): + self.id = id + self.replacement = replacement + super().__init__() + + def visit_TargetHole(self, node: TargetHole) -> ast.expr: + return self.replacement if node.id == self.id else node + + +def _bind_target_hole( + stack: list[ast.expr], hole: ast.expr, replacement: ast.expr +) -> list[ast.expr]: + """Fill the loop-target hole ``hole`` of the innermost matching comprehension. + + Returns a new stack; the hole itself is left in place for the caller to pop. + """ + assert isinstance(hole, TargetHole), f"Expected a loop target hole, got {hole}" + for pos, item in zip(reversed(range(len(stack))), reversed(stack)): + if not isinstance(item, CompExp) or not item.generators: + continue + if not any( + isinstance(n, TargetHole) and n.id == hole.id + for n in ast.walk(item.generators[-1].target) + ): + continue + new_comp = ReplaceTargetHole(hole.id, replacement).visit(copy.deepcopy(item)) + return stack[:pos] + [new_comp] + stack[pos + 1 :] + + raise TypeError(f"No comprehension found with loop target hole {hole.id}") + + +class Null(ast.Constant): + """Placeholder for NULL values generated in bytecode.""" + + def __init__(self, value: None = None): + super().__init__(value=value) + + +class ConvertedValue(ast.expr): + """Wrapper for values that have been converted with CONVERT_VALUE.""" + + value: ast.expr + conversion: int + ast_conversion: int + + def __init__(self, value: ast.expr, conversion: int): + self.value = value + self.conversion = conversion + # Map CONVERT_VALUE args to ast.FormattedValue conversion values + # CONVERT_VALUE: 0=None, 1=str, 2=repr, 3=ascii + # ast.FormattedValue: -1=none, 115=str, 114=repr, 97=ascii + conversion_map = {0: -1, 1: 115, 2: 114, 3: 97} + self.ast_conversion = conversion_map.get(conversion, -1) + + +class CompLambda(ast.Lambda): + """Placeholder AST node representing a lambda function used in comprehensions.""" + + def __init__(self, body: CompExp): + assert isinstance(body, CompExp) + assert sum(1 for x in ast.walk(body) if isinstance(x, DummyIterName)) == 1 + assert len(body.generators) > 0 + assert isinstance(body.generators[0].iter, DummyIterName) + args = ast.arguments( + posonlyargs=[ast.arg(DummyIterName().id)], + args=[], + kwonlyargs=[], + kw_defaults=[], + defaults=[], + ) + super().__init__(args=args, body=body) + + def __copy__(self): + """Support copy.copy operation.""" + assert isinstance(self.body, CompExp) + return CompLambda(self.body) + + def __deepcopy__(self, memo): + """Support copy.deepcopy operation.""" + assert isinstance(self.body, CompExp) + return CompLambda(copy.deepcopy(self.body, memo)) + + def inline(self, iterator: ast.expr) -> CompExp: + assert isinstance(self.body, CompExp) + res: CompExp = copy.deepcopy(self.body) + res.generators[0].iter = iterator + return res + + +class ReplacePlaceholder(ast.NodeTransformer): + value: ast.expr + _done: bool + + def __init__(self, value: ast.expr): + self.value = value + self._done = False + super().__init__() + + def visit(self, node): + if isinstance(node, Placeholder) and not self._done: + self._done = True + return self.value + else: + return self.generic_visit(node) + + +class ReplaceSkipped(ast.NodeTransformer): + id: str + replacement: ast.expr + + def __init__(self, id: str, replacement: ast.expr): + self.id = id + self.replacement = copy.deepcopy(replacement) + super().__init__() + + def visit_IfExp(self, node: ast.IfExp): + if isinstance(node.body, Skipped) and node.body.id == self.id: + return ast.IfExp(test=node.test, body=self.replacement, orelse=node.orelse) + elif isinstance(node.orelse, Skipped) and node.orelse.id == self.id: + return ast.IfExp(test=node.test, body=node.body, orelse=self.replacement) + else: + return self.generic_visit(node) + + +class BranchState(typing.NamedTuple): + testval: bool + value: ast.expr + + +class BranchIdentifier(ast.NodeVisitor): + branching: collections.abc.MutableMapping[str, BranchState] + filter_lengths: list[int] + + def __init__(self): + self.branching = {} + self.filter_lengths = [] + super().__init__() + + def visit_IfExp(self, node: ast.IfExp): + if isinstance(node.body, Skipped): + self.branching[node.body.id] = BranchState( + testval=False, value=copy.deepcopy(node.orelse) + ) + elif isinstance(node.orelse, Skipped): + self.branching[node.orelse.id] = BranchState( + testval=True, value=copy.deepcopy(node.body) + ) + return self.generic_visit(node) + + def visit_comprehension(self, node: ast.comprehension): + self.filter_lengths.append(len(node.ifs)) + return self.generic_visit(node) + + +@functools.cache +def _instructions( + code: types.CodeType, +) -> collections.abc.Mapping[int, dis.Instruction]: + """Decode a code object once; every state derived from it shares the result.""" + return collections.OrderedDict( + (instr.offset, instr) for instr in dis.get_instructions(code) + ) + + +@functools.cache +def _next_instructions( + code: types.CodeType, +) -> collections.abc.Mapping[int, dis.Instruction]: + """Map each instruction offset to the instruction that follows it.""" + ordered = list(_instructions(code).values()) + return {before.offset: after for before, after in zip(ordered[:-1], ordered[1:])} + + +@dataclass(frozen=True) +class ReconstructionState: + """State maintained during AST reconstruction from bytecode. + + This class tracks all the information needed while processing bytecode + instructions to reconstruct the original comprehension's AST. It acts + as the working memory during the reconstruction process, maintaining + both the evaluation stack state and the high-level comprehension structure + being built. + + The reconstruction process works by simulating the Python VM's execution + of the bytecode, but instead of executing operations, it builds AST nodes + that represent those operations. + + Attributes: + code: The compiled code object from which the bytecode is being processed. + This is typically obtained from a generator function or comprehension. + + stack: Simulates the Python VM's value stack. Contains AST nodes or + values that would be on the stack during execution. Operations + like LOAD_FAST push to this stack, while operations like + BINARY_ADD pop operands and push results. + """ + + code: types.CodeType + instruction: dis.Instruction + + stack: list[ast.expr] = field(default_factory=list) + result: ast.expr = field(default_factory=Placeholder) + + # How many times each FOR_ITER has been entered on this path. + loops: dict[int, int] = field(default_factory=dict) + finished: bool = field(default=False) + + # Which edge each already-resolved conditional jump took on this path. + branches: "dict[int, BranchEdge]" = field(default_factory=dict) + + # Locals bound to a known expression rather than to a loop target. 3.14 + # unrolls a single-iteration loop over a literal, storing its targets + # directly, so those names have to be substituted back at their uses. + bindings: dict[str, ast.expr] = field(default_factory=dict) + + # Set by KW_NAMES (Python 3.12 only) and consumed by the following CALL. + # KW_NAMES has no stack effect, so the names cannot live on `stack`. + kw_names: tuple[str, ...] | None = field(default=None) + + # The value captured in each closure cell the code reads, by name. A free + # variable is looked up in a cell, not in globals, so reconstructing it as a + # bare name would resolve to whatever the evaluating namespace happens to + # bind; the captured value has to be written into the tree instead. + freevars: dict[str, ast.expr] = field(default_factory=dict) + + @property + def instructions(self) -> collections.abc.Mapping[int, dis.Instruction]: + """The bytecode instructions of the current code object, by offset.""" + return _instructions(self.code) + + @property + def next_instructions(self) -> collections.abc.Mapping[int, dis.Instruction]: + return _next_instructions(self.code) + + +# Python version enum for version-specific handling +class PythonVersion(enum.IntEnum): + PY_312 = 12 + PY_313 = 13 + PY_314 = 14 + + +def current_version() -> PythonVersion: + """The bytecode dialect of the running interpreter. + + Raises on a Python this module has not been taught, rather than guessing + that the previous release's opcodes still mean what they used to. + """ + try: + return PythonVersion(sys.version_info.minor) + except ValueError as e: + supported = ", ".join(f"3.{v.value}" for v in PythonVersion) + raise NotImplementedError( + f"effectful.internals.disassembly supports {supported}, " + f"not 3.{sys.version_info.minor}" + ) from e + + +# Global handler registry +OpHandler = Callable[[ReconstructionState, dis.Instruction], ReconstructionState] + +OP_HANDLERS: dict[str, OpHandler] = {} + + +@typing.overload +def register_handler( + opname: str, *, version: PythonVersion +) -> Callable[[OpHandler], OpHandler]: ... + + +@typing.overload +def register_handler( + opname: str, + handler: OpHandler, + *, + version: PythonVersion, +) -> OpHandler: ... + + +def register_handler( + opname: str, + handler=None, + *, + version: PythonVersion, +): + """Register a handler for one opcode in one Python bytecode dialect. + + Every dialect a handler applies to is named explicitly. Opcodes are not + assumed to carry forward: a release can keep an opcode's name while changing + what it does, so applicability to a new Python is a decision to make per + opcode rather than a default. + """ + if handler is None: + return functools.partial(register_handler, opname, version=version) + + # Skip registration if version doesn't match current version + if version != current_version(): + return handler + + # Only check opmap if the version matches (or no version specified) + assert opname in dis.opmap, f"Invalid operation name: '{opname}'" + + if opname in OP_HANDLERS: + raise ValueError(f"Handler for '{opname}' (version {version}) already exists.") + + if dis.opmap[opname] in dis.hasjrel: + assert opname in LOOP_OPS | BRANCH_OPS | JUMP_OPS + else: + assert opname not in LOOP_OPS | BRANCH_OPS | JUMP_OPS + + @functools.wraps(handler) + def _wrapper( + state: ReconstructionState, + instr: dis.Instruction, + ) -> ReconstructionState: + assert instr.opname == opname, ( + f"Handler for '{opname}' called with wrong instruction" + ) + assert not state.finished, "Cannot process instruction on finished state" + + new_state = handler(state, instr) + + jump: bool | None # argument to dis.stack_effect + if instr.opname in LOOP_OPS: + if state.loops.get(instr.offset, 0) > 0: + new_state = replace( + new_state, instruction=state.instructions[instr.argval] + ) + jump = True + else: + # Copy rather than mutate: continuations forked from this state + # share the mapping and must not see each other's loop counts. + new_state = replace( + new_state, + instruction=state.next_instructions[instr.offset], + loops={ + **state.loops, + instr.offset: state.loops.get(instr.offset, 0) + 1, + }, + ) + jump = False + elif instr.opname in BRANCH_OPS: + if new_state.branches.get(instr.offset) == BranchEdge.FALL_THROUGH: + new_state = replace( + new_state, instruction=state.next_instructions[instr.offset] + ) + jump = False + else: + new_state = replace( + new_state, instruction=state.instructions[instr.argval] + ) + jump = True + elif instr.opname in JUMP_OPS: + new_state = replace(new_state, instruction=state.instructions[instr.argval]) + jump = True + elif instr.opname not in RETURN_OPS and instr.offset in state.next_instructions: + new_state = replace( + new_state, instruction=state.next_instructions[instr.offset] + ) + jump = None + else: + new_state = replace(new_state, finished=True) + jump = None + + # post-condition: check stack effect + expected_stack_effect = dis.stack_effect(instr.opcode, instr.arg, jump=jump) + actual_stack_effect = len(new_state.stack) - len(state.stack) + assert len(state.stack) + expected_stack_effect >= 0, ( + f"Handler for '{opname}' would result in negative stack size" + ) + assert actual_stack_effect == expected_stack_effect, ( + f"Handler for '{opname}' has incorrect stack effect: " + f"expected {expected_stack_effect}, got {actual_stack_effect}" + ) + + return new_state + + OP_HANDLERS[opname] = _wrapper + return handler # return the original handler for multiple decorator usage + + +LOOP_OPS: set[typing.Literal["FOR_ITER"]] = {"FOR_ITER"} + +BRANCH_OPS: set[ + typing.Literal[ + "POP_JUMP_IF_TRUE", + "POP_JUMP_IF_FALSE", + "POP_JUMP_IF_NOT_NONE", + "POP_JUMP_IF_NONE", + ] +] = { + "POP_JUMP_IF_TRUE", + "POP_JUMP_IF_FALSE", + "POP_JUMP_IF_NOT_NONE", + "POP_JUMP_IF_NONE", +} + +RETURN_OPS: set[typing.Literal["RETURN_VALUE", "RETURN_CONST"]] = { + "RETURN_VALUE", + "RETURN_CONST", +} + +JUMP_OPS = {dis.opname[d] for d in dis.hasjrel} - LOOP_OPS - BRANCH_OPS - RETURN_OPS + + +# Instructions that emit an element of the comprehension being built. Reaching +# one of these means the current iteration was *not* filtered out. +PRODUCE_OPS = {"YIELD_VALUE", "LIST_APPEND", "SET_ADD", "MAP_ADD"} + + +def _successor_offsets(state: ReconstructionState, instr: dis.Instruction) -> list[int]: + """Offsets control can transfer to from ``instr``, ignoring exception edges.""" + following = state.next_instructions.get(instr.offset) + if instr.opname in BRANCH_OPS | LOOP_OPS: + return [instr.argval] + ([following.offset] if following else []) + elif instr.opname in JUMP_OPS: + return [instr.argval] + elif instr.opname in RETURN_OPS: + return [] + else: + return [following.offset] if following else [] + + +def _reachable_outcomes(state: ReconstructionState, start: int) -> tuple[bool, bool]: + """From ``start``, can the iteration be skipped, and can an element be produced? + + Returns ``(can_skip, can_produce)``. "Skip" means reaching the loop + back-edge without emitting an element, i.e. being filtered out. + """ + seen: set[int] = set() + pending = [start] + can_skip = can_produce = False + + while pending: + offset = pending.pop() + if offset in seen or offset not in state.instructions: + continue + seen.add(offset) + + instr = state.instructions[offset] + if instr.opname in PRODUCE_OPS: + can_produce = True + continue + if ( + instr.opname == "JUMP_BACKWARD" + and state.instructions[instr.argval].opname in LOOP_OPS + ): + can_skip = True + continue + + pending.extend(_successor_offsets(state, instr)) + + return can_skip, can_produce + + +class BranchEdge(enum.IntEnum): + """Which way a conditional jump was resolved on the path being explored.""" + + TAKE_JUMP = 1 + FALL_THROUGH = 2 + + +class BranchKind(enum.Enum): + """What role a conditional jump plays in a comprehension. + + TERNARY + A conditional expression: the arms reconverge having each pushed a + value, and both are spliced back together into an ``ast.IfExp``. + FILTER + Part of a filter's condition. The condition is consumed rather than + producing a value, so each surviving path records the conjunction of + tests that got it to the element, and the filter as a whole is the + disjunction of those conjunctions. + """ + + TERNARY = enum.auto() + FILTER = enum.auto() + + +@functools.cache +def _stack_depths(code: types.CodeType) -> collections.abc.Mapping[int, int]: + """VM stack depth on entry to each reachable instruction. + + Depths are relative to the start of the code object, which is all that is + needed to tell a value-producing branch from a control-flow one. + """ + instructions = collections.OrderedDict( + (i.offset, i) for i in dis.get_instructions(code) + ) + ordered = list(instructions.values()) + following = {a.offset: b.offset for a, b in zip(ordered[:-1], ordered[1:])} + + depths: dict[int, int] = {ordered[0].offset: 0} + pending = collections.deque([ordered[0].offset]) + while pending: + offset = pending.popleft() + instr, depth = instructions[offset], depths[offset] + if instr.opname in RETURN_OPS: + continue + + edges: list[tuple[int, bool | None]] = [] + if instr.opname in BRANCH_OPS | LOOP_OPS: + edges = [(instr.argval, True)] + if offset in following: + edges.append((following[offset], False)) + elif instr.opname in JUMP_OPS: + edges = [(instr.argval, True)] + elif offset in following: + edges = [(following[offset], None)] + + for target, jump in edges: + if target in instructions and target not in depths: + depths[target] = depth + dis.stack_effect( + instr.opcode, instr.arg, jump=jump + ) + pending.append(target) + + return depths + + +def _forward_reachable(state: ReconstructionState, start: int) -> set[int]: + """Offsets reachable from ``start`` without producing or looping back.""" + seen: set[int] = set() + pending = [start] + while pending: + offset = pending.pop() + if offset in seen or offset not in state.instructions: + continue + seen.add(offset) + + instr = state.instructions[offset] + if instr.opname in PRODUCE_OPS: + continue + if ( + instr.opname == "JUMP_BACKWARD" + and state.instructions[instr.argval].opname in LOOP_OPS + ): + continue + + pending.extend(_successor_offsets(state, instr)) + + return seen + + +def _is_conditional_expression( + state: ReconstructionState, instr: dis.Instruction +) -> bool: + """Do the two edges of ``instr`` reconverge one stack slot deeper? + + That is the signature of a conditional expression: each arm leaves a value + behind and control rejoins to consume it. A filter's condition is consumed + by the jump itself, so wherever its edges meet -- if they meet at all -- the + stack is no deeper than it was. + """ + following = state.next_instructions.get(instr.offset) + if following is None: + return False + + common = _forward_reachable(state, instr.argval) & _forward_reachable( + state, following.offset + ) + if not common: + return False # the edges never rejoin, so nothing was left on the stack + + depths = _stack_depths(state.code) + join = min(common) # arms are laid out contiguously, so the join comes first + if join not in depths or following.offset not in depths: + return False + return depths[join] == depths[following.offset] + 1 + + +def _classify_branch( + state: ReconstructionState, instr: dis.Instruction +) -> tuple[BranchKind, list[BranchEdge]]: + """Classify a conditional jump and list the edges worth exploring.""" + both = [BranchEdge.TAKE_JUMP, BranchEdge.FALL_THROUGH] + following = state.next_instructions.get(instr.offset) + if following is None: + return BranchKind.TERNARY, both + + jump_skip, _ = _reachable_outcomes(state, instr.argval) + fall_skip, _ = _reachable_outcomes(state, following.offset) + + # Neither edge can drop the current iteration -- because there is no loop at + # all (a lambda body) or because the element is produced regardless (a + # conditional in the element expression). Either way nothing is filtered. + if not jump_skip and not fall_skip: + return BranchKind.TERNARY, both + + # Otherwise the branch could be a filter, or a conditional expression that + # merely happens to sit inside one. Only the latter leaves a value behind. + if _is_conditional_expression(state, instr): + return BranchKind.TERNARY, both + + # An edge that cannot reach an element contributes nothing to the filter, so + # there is no point walking it. Pruning those edges is also what keeps the + # executor out of the operand-cleanup block on a chained comparison's + # failing edge. + live = [ + edge + for edge, start in ( + (BranchEdge.TAKE_JUMP, instr.argval), + (BranchEdge.FALL_THROUGH, following.offset), + ) + if _reachable_outcomes(state, start)[1] + ] + return BranchKind.FILTER, live or [BranchEdge.TAKE_JUMP] + + +def _negate(condition: ast.expr) -> ast.expr: + """Logical negation, cancelling a `not` rather than stacking another one.""" + if isinstance(condition, ast.UnaryOp) and isinstance(condition.op, ast.Not): + return condition.operand + return ast.UnaryOp(op=ast.Not(), operand=condition) + + +def _conjoin(conditions: list[ast.expr]) -> ast.expr | None: + """Combine the entries of a ``comprehension.ifs`` list into one expression.""" + if not conditions: + return None + elif len(conditions) == 1: + return conditions[0] + else: + return ast.BoolOp(op=ast.And(), values=list(conditions)) + + +def _disjoin(left: ast.expr, right: ast.expr) -> ast.expr: + """Combine two conditions with ``or``, flattening nested disjunctions. + + Two rewrites are applied while combining. Duplicate disjuncts are dropped, + because paths through independent filters repeat them. And ``X or (not X and + Y)`` becomes ``X or Y``: enumerating paths records the negation of every + test a path declined, so a later disjunct restates the negation of an + earlier one. Dropping it is not merely tidier -- `or` short-circuits, so the + earlier disjunct has already been evaluated, and leaving the negation in + would evaluate it a second time, which is visibly wrong when the condition + contains an assignment expression. + + Conditions are keyed by ``ast.dump`` exactly once each: these lists get long + and the expressions large, so re-dumping per comparison dominates. + """ + values: list[ast.expr] = [] + for side in (left, right): + if isinstance(side, ast.BoolOp) and isinstance(side.op, ast.Or): + values.extend(side.values) + else: + values.append(side) + + unique: list[ast.expr] = [] + seen: set[str] = set() + for value in values: + key = ast.dump(value) + if key in seen: + continue + seen.add(key) + + # Absorb the negations of the disjuncts already accepted. + if isinstance(value, ast.BoolOp) and isinstance(value.op, ast.And): + kept = [ + conjunct + for conjunct in value.values + if not ( + isinstance(conjunct, ast.UnaryOp) + and isinstance(conjunct.op, ast.Not) + and ast.dump(conjunct.operand) in seen + ) + ] + if kept and len(kept) < len(value.values): + conjoined = _conjoin(kept) + assert conjoined is not None + value = conjoined + + unique.append(value) + + return unique[0] if len(unique) == 1 else ast.BoolOp(op=ast.Or(), values=unique) + + +def _merge_filters_into( + node: typing.Any, other: typing.Any, mutate: bool = True +) -> bool: + """Walk two results in parallel, OR-ing the filters where they disagree. + + Everything outside a ``comprehension.ifs`` has to match exactly; the ifs are + where the paths are allowed to differ, and are combined rather than + compared. Filters are not recursed into, so a nested comprehension inside a + filter is treated as part of that filter's condition. + + Returns False if the two results differ somewhere they may not. With + ``mutate=False`` nothing is written, which allows compatibility to be tested + before paying for a deep copy -- most candidate pairs do not merge, and the + copy dominates otherwise. + """ + if type(node) is not type(other): + return False + + if isinstance(node, ast.comprehension): + if ast.dump(node.target) != ast.dump(other.target): + return False + if not _merge_filters_into(node.iter, other.iter, mutate): + return False + if not mutate: + return True + + guard, other_guard = _conjoin(node.ifs), _conjoin(other.ifs) + if guard is None or other_guard is None: + # One path reached the element unconditionally, so the filter as a + # whole is unconditional at this generator. + node.ifs = [] + elif ast.dump(guard) != ast.dump(other_guard): + node.ifs = [_disjoin(guard, other_guard)] + return True + + if not isinstance(node, ast.AST): + return bool(node == other) + + for name in node._fields: + mine, theirs = getattr(node, name, None), getattr(other, name, None) + if isinstance(mine, list) or isinstance(theirs, list): + if not isinstance(mine, list) or not isinstance(theirs, list): + return False + if len(mine) != len(theirs): + return False + if not all(_merge_filters_into(a, b, mutate) for a, b in zip(mine, theirs)): + return False + elif isinstance(mine, ast.AST) or isinstance(theirs, ast.AST): + if not isinstance(mine, ast.AST) or not isinstance(theirs, ast.AST): + return False + if not _merge_filters_into(mine, theirs, mutate): + return False + elif mine != theirs: + return False + + return True + + +def _merge_filters(left: ast.expr, right: ast.expr) -> ast.expr | None: + """Combine two paths that differ only in which filter conditions they met. + + Each path through a filter records the conjunction that got it to the + element; the filter as a whole is the disjunction over all such paths. + Returns ``None`` when the results differ by more than their filters. + """ + # A marker anywhere means some conditional expression is still unresolved, + # and unresolved arms must be spliced before anything can be OR-ed. + if any(isinstance(n, Skipped) for n in ast.walk(left)): + return None + if any(isinstance(n, Skipped) for n in ast.walk(right)): + return None + + if not _merge_filters_into(left, right, mutate=False): + return None + + merged = copy.deepcopy(left) + return merged if _merge_filters_into(merged, right) else None + + +def _skipped_offset(key: str) -> int: + """Sort key for `.SKIPPED_` markers, so merging is deterministic.""" + return int(key.rsplit("_", 1)[-1]) + + +def _merge_at_ifexp(left: ast.expr, right: ast.expr) -> ast.expr: + """ + Merge two expression ASTs obtained from two branches of symbolic execution. + """ + if isinstance(left, ast.Constant) and left.value is None: + return copy.deepcopy(right) + elif isinstance(right, ast.Constant) and right.value is None: + return copy.deepcopy(left) + + assert type(left) == type(right) + + lb, rb = BranchIdentifier(), BranchIdentifier() + lb.visit(left) + rb.visit(right) + + # A conditional expression: each path filled in one arm and left a marker in + # the other, so splice the two together. Sorted for determinism -- set + # iteration order over the marker names varies with PYTHONHASHSEED. + common_keys = set(lb.branching) & set(rb.branching) + differing = [ + key + for key in sorted(common_keys, key=_skipped_offset) + if lb.branching[key].testval != rb.branching[key].testval + ] + + # Only copy once it is known there is something to splice; this runs for + # every candidate pair of paths, most of which have nothing in common. + merged: ast.expr = copy.deepcopy(left) if differing else left + for key in differing: + visited = ReplaceSkipped(key, rb.branching[key].value).visit(merged) + assert isinstance(visited, ast.expr) + merged = visited + spliced = bool(differing) + + # The paths may *also* have satisfied different filter conditions on the way + # to the element, so combine those too rather than picking one arbitrarily. + combined = _merge_filters(merged, right) + if combined is not None: + return combined + if spliced: + return merged + + if ast.dump(left) == ast.dump(right): + return copy.deepcopy(left) + + raise ValueError("No differing branches found to merge") + + +def _specialization_guard_edge( + state: ReconstructionState, instr: dis.Instruction +) -> BranchEdge | None: + """The edge past a 3.14 inlined-builtin guard, or None if this isn't one. + + 3.14 may inline `any()`/`all()` over a generator, guarding the inlined code + with `loaded_name is ` and keeping an ordinary call on the + other edge. Only that other edge still contains the call to reconstruct, so + the guard is treated as though the identity test failed. + """ + if not state.stack: + return None + condition = state.stack[-1] + if not isinstance(condition, ast.Compare): + return None + if not any(isinstance(c, CommonConstant) for c in condition.comparators): + return None + + # Follow the edge taken when the identity test is false. + if instr.opname == "POP_JUMP_IF_FALSE": + return BranchEdge.TAKE_JUMP + elif instr.opname == "POP_JUMP_IF_TRUE": + return BranchEdge.FALL_THROUGH + else: + return None + + +def _merge_all(results: list[ast.expr]) -> ast.expr: + """Combine every path's result into one expression. + + Merging is not associative: a path that still carries an unfilled + conditional-expression arm can only combine with the path that took the + other arm, which need not be its neighbour. So rather than folding in + order, repeatedly merge whichever pair actually combines. + """ + pending = list(results) + while len(pending) > 1: + for i, j in itertools.combinations(range(len(pending)), 2): + try: + merged = _merge_at_ifexp(pending[i], pending[j]) + except (ValueError, AssertionError): + continue + pending = [merged] + [p for k, p in enumerate(pending) if k not in (i, j)] + break + else: + raise ValueError("Could not merge the paths of symbolic execution") + + return pending[0] + + +def _decide_branch( + state: ReconstructionState, instr: dis.Instruction, edge: BranchEdge +) -> ReconstructionState: + """Record which edge of ``instr`` the path being explored takes.""" + return replace(state, branches={**state.branches, instr.offset: edge}) + + +def _symbolic_exec(code: types.CodeType, freevars: dict[str, ast.expr]) -> ast.expr: + """Execute bytecode symbolically, following control flow.""" + continuations: list[ReconstructionState] = [ + ReconstructionState( + code=code, + instruction=next(iter(dis.get_instructions(code))), + stack=[Placeholder(), Placeholder()] + if current_version() == PythonVersion.PY_312 + and code.co_flags & inspect.CO_GENERATOR + else [Placeholder()], + freevars=freevars, + ) + ] + + results: list[ast.expr] = [] + + while continuations: + state = continuations.pop() + while not state.finished: + instr = state.instruction + if instr.opname in BRANCH_OPS and instr.offset not in state.branches: + forced = _specialization_guard_edge(state, instr) + if forced is not None: + state = _decide_branch(state, instr, forced) + else: + _, live = _classify_branch(state, instr) + # Explore the first live edge now; queue the rest for later. + continuations.extend( + _decide_branch(state, instr, edge) for edge in live[1:] + ) + state = _decide_branch(state, instr, live[0]) + + state = OP_HANDLERS[state.instruction.opname](state, state.instruction) + results.append(state.result) + + assert results, "No results from symbolic execution" + result = _merge_all(results) + assert not any(isinstance(n, Skipped) for n in ast.walk(result)), ( + "Every conditional expression arm must have been filled in" + ) + return result + + +# ============================================================================ +# GENERATOR COMPREHENSION HANDLERS +# ============================================================================ + + +@register_handler("RETURN_GENERATOR", version=PythonVersion.PY_312) +def handle_return_generator_312( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # RETURN_GENERATOR is the first instruction in generator expressions in Python 3.13+ + assert len(state.stack) == 2 and all( + isinstance(x, Null | Placeholder) for x in state.stack + ), "RETURN_GENERATOR must be the first instruction" + new_result = ast.GeneratorExp(elt=Placeholder(), generators=[]) + return replace(state, stack=[new_result, Null()]) + + +@register_handler("RETURN_GENERATOR", version=PythonVersion.PY_313) +@register_handler("RETURN_GENERATOR", version=PythonVersion.PY_314) +def handle_return_generator( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # RETURN_GENERATOR is the first instruction in generator expressions in Python 3.13+ + assert len(state.stack) == 1 and isinstance(state.stack[0], Null | Placeholder), ( + "RETURN_GENERATOR must be the first instruction" + ) + return replace( + state, stack=[ast.GeneratorExp(elt=Placeholder(), generators=[]), Null()] + ) + + +@register_handler("YIELD_VALUE", version=PythonVersion.PY_312) +@register_handler("YIELD_VALUE", version=PythonVersion.PY_313) +@register_handler("YIELD_VALUE", version=PythonVersion.PY_314) +def handle_yield_value( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # YIELD_VALUE pops a value from the stack and yields it + # This is the expression part of the generator + assert isinstance(state.result, Placeholder) + new_result = copy.deepcopy(state.stack[0]) + assert isinstance(new_result, ast.GeneratorExp), ( + "YIELD_VALUE must be called after RETURN_GENERATOR" + ) + assert len(new_result.generators) > 0, "YIELD_VALUE should have generators" + assert any(isinstance(x, Placeholder) for x in ast.walk(new_result.elt)) + new_result.elt = ReplacePlaceholder(ensure_ast(state.stack[-1])).visit( + new_result.elt + ) + new_stack = [new_result] + state.stack[1:] + return replace(state, stack=new_stack, result=new_result) + + +# ============================================================================ +# LIST COMPREHENSION HANDLERS +# ============================================================================ + + +@register_handler("BUILD_LIST", version=PythonVersion.PY_312) +@register_handler("BUILD_LIST", version=PythonVersion.PY_313) +@register_handler("BUILD_LIST", version=PythonVersion.PY_314) +def handle_build_list( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + assert instr.arg is not None + size: int = instr.arg + + if size == 0: + # Check if this looks like the start of a list comprehension pattern + # In nested comprehensions, BUILD_LIST(0) starts a new list comprehe + new_ret = ast.ListComp(elt=Placeholder(), generators=[]) + new_stack = state.stack + [new_ret] + return replace(state, stack=new_stack) + else: + # BUILD_LIST with elements - create a regular list + elements = [ensure_ast(elem) for elem in state.stack[-size:]] + new_stack = state.stack[:-size] + elt_node = ast.List(elts=elements, ctx=ast.Load()) + new_stack = new_stack + [elt_node] + return replace(state, stack=new_stack) + + +@register_handler("LIST_APPEND", version=PythonVersion.PY_312) +@register_handler("LIST_APPEND", version=PythonVersion.PY_313) +@register_handler("LIST_APPEND", version=PythonVersion.PY_314) +def handle_list_append( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + assert isinstance(state.stack[-instr.argval - 1], ast.ListComp) + + # add the body to the comprehension + comp: ast.ListComp = copy.deepcopy(state.stack[-instr.argval - 1]) + assert any(isinstance(x, Placeholder) for x in ast.walk(comp.elt)) + comp.elt = ReplacePlaceholder(state.stack[-1]).visit(comp.elt) + + # swap the return value + new_stack = state.stack[:-1] + new_stack[-instr.argval] = comp + + return replace(state, stack=new_stack) + + +# ============================================================================ +# SET COMPREHENSION HANDLERS +# ============================================================================ + + +@register_handler("BUILD_SET", version=PythonVersion.PY_312) +@register_handler("BUILD_SET", version=PythonVersion.PY_313) +@register_handler("BUILD_SET", version=PythonVersion.PY_314) +def handle_build_set( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + assert instr.arg is not None + size: int = instr.arg + + if size == 0: + new_result = ast.SetComp(elt=Placeholder(), generators=[]) + new_stack = state.stack + [new_result] + return replace(state, stack=new_stack) + else: + elements = [ensure_ast(elem) for elem in state.stack[-size:]] + new_stack = state.stack[:-size] + elt_node = ast.Set(elts=elements) + new_stack = new_stack + [elt_node] + return replace(state, stack=new_stack) + + +@register_handler("SET_ADD", version=PythonVersion.PY_312) +@register_handler("SET_ADD", version=PythonVersion.PY_313) +@register_handler("SET_ADD", version=PythonVersion.PY_314) +def handle_set_add( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + assert isinstance(state.stack[-instr.argval - 1], ast.SetComp) + + # add the body to the comprehension + comp: ast.SetComp = copy.deepcopy(state.stack[-instr.argval - 1]) + assert any(isinstance(x, Placeholder) for x in ast.walk(comp.elt)) + comp.elt = ReplacePlaceholder(state.stack[-1]).visit(comp.elt) + + # swap the return value + new_stack = state.stack[:-1] + new_stack[-instr.argval] = comp + + return replace(state, stack=new_stack) + + +# ============================================================================ +# DICT COMPREHENSION HANDLERS +# ============================================================================ + + +@register_handler("BUILD_MAP", version=PythonVersion.PY_312) +@register_handler("BUILD_MAP", version=PythonVersion.PY_313) +@register_handler("BUILD_MAP", version=PythonVersion.PY_314) +def handle_build_map( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + assert instr.arg is not None + size: int = instr.arg + + if size == 0: + new_result = ast.DictComp(key=Placeholder(), value=Placeholder(), generators=[]) + new_stack = state.stack + [new_result] + return replace(state, stack=new_stack) + else: + # Pop key-value pairs for the dict. They sit on the stack in source + # order -- key first, then value -- so they must be read from the + # bottom of that slice up: a later duplicate key has to keep winning, + # and side effects have to happen in the order they were written. + pairs = state.stack[-2 * size :] + keys: list[ast.expr | None] = [ensure_ast(pairs[2 * i]) for i in range(size)] + values = [ensure_ast(pairs[2 * i + 1]) for i in range(size)] + new_stack = state.stack[: -2 * size] + + # Create dict AST + dict_node = ast.Dict(keys=keys, values=values) + new_stack = new_stack + [dict_node] + return replace(state, stack=new_stack) + + +@register_handler("MAP_ADD", version=PythonVersion.PY_312) +@register_handler("MAP_ADD", version=PythonVersion.PY_313) +@register_handler("MAP_ADD", version=PythonVersion.PY_314) +def handle_map_add( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + assert isinstance(state.stack[-instr.argval - 2], ast.DictComp) + + # add the body to the comprehension + comp: ast.DictComp = copy.deepcopy(state.stack[-instr.argval - 2]) + assert any(isinstance(x, Placeholder) for x in ast.walk(comp.key)) + assert any(isinstance(x, Placeholder) for x in ast.walk(comp.value)) + comp.key = ReplacePlaceholder(state.stack[-2]).visit(comp.key) + comp.value = ReplacePlaceholder(state.stack[-1]).visit(comp.value) + + # swap the return value + new_stack = state.stack[:-2] + new_stack[-instr.argval] = comp + + return replace(state, stack=new_stack) + + +# ============================================================================ +# LOOP CONTROL HANDLERS +# ============================================================================ + + +@register_handler("RETURN_VALUE", version=PythonVersion.PY_312) +@register_handler("RETURN_VALUE", version=PythonVersion.PY_313) +def handle_return_value( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + assert isinstance(state.result, Placeholder) + assert len(state.stack) == 2 + new_result = ReplacePlaceholder(ensure_ast(state.stack[-1])).visit(state.stack[-2]) + new_stack = state.stack[:-1] + return replace(state, stack=new_stack, result=new_result) + + +def _unyielded_comprehension(state: ReconstructionState) -> CompExp | None: + """The comprehension of a body the compiler proved unreachable, if any. + + An always-false filter lets the compiler drop the whole body: the loop is + still walked, but nothing is ever yielded or appended, so no element + expression survives. The partly built comprehension still carries its + generators, so it can be rebuilt with a filter that is never satisfied -- + which iterates exactly as the original did and produces nothing. + """ + for item in reversed(state.stack): + if not isinstance(item, CompExp) or not item.generators: + continue + + element = item.value if isinstance(item, ast.DictComp) else item.elt + if not isinstance(element, Placeholder): + continue + + unreachable = copy.deepcopy(item) + never = ast.Constant(value=None) + if isinstance(unreachable, ast.DictComp): + unreachable.key, unreachable.value = never, copy.deepcopy(never) + else: + unreachable.elt = never + unreachable.generators[-1].ifs = [ast.Constant(value=False)] + return unreachable + + return None + + +@register_handler("RETURN_VALUE", version=PythonVersion.PY_314) +def handle_return_value_314( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # Two things changed in 3.14. RETURN_CONST is gone, so a generator's + # trailing `return None` now arrives as LOAD_CONST + RETURN_VALUE; and + # RETURN_VALUE's stack effect is 0 rather than -1, the returned value being + # discarded along with the frame. The value therefore stays on the stack. + if not isinstance(state.result, Placeholder): + assert ( + isinstance(state.stack[-1], ast.Constant) and state.stack[-1].value is None + ), "A generator may only fall off the end returning None" + return state + + unreachable = _unyielded_comprehension(state) + if unreachable is not None: + return replace(state, result=unreachable) + + assert len(state.stack) == 2 + new_result = ReplacePlaceholder(ensure_ast(state.stack[-1])).visit(state.stack[-2]) + return replace(state, result=new_result) + + +@register_handler("RETURN_CONST", version=PythonVersion.PY_312) +@register_handler("RETURN_CONST", version=PythonVersion.PY_313) +def handle_return_const( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # RETURN_CONST returns a constant value (replaces some LOAD_CONST + RETURN_VALUE patterns) + # Similar to RETURN_VALUE but with a constant + if isinstance(state.result, Placeholder): + unreachable = _unyielded_comprehension(state) + if unreachable is not None: + return replace(state, result=unreachable) + return replace(state, result=ensure_ast(instr.argval)) + else: + assert instr.argval is None + return state + + +@register_handler("FOR_ITER", version=PythonVersion.PY_312) +@register_handler("FOR_ITER", version=PythonVersion.PY_313) +@register_handler("FOR_ITER", version=PythonVersion.PY_314) +def handle_for_iter( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # FOR_ITER pops an iterator from the stack and pushes the next item + # If the iterator is exhausted, it jumps to the target instruction + assert len(state.stack) > 0, "FOR_ITER must have an iterator on the stack" + + if state.loops.get(instr.offset, 0) > 0: + return replace(state, stack=state.stack + [Null()]) + + # The iterator should be on top of stack + iterator: ast.expr = state.stack[-1] + + for pos, item in zip(reversed(range(len(state.stack))), reversed(state.stack)): + if not isinstance(item, CompExp): + continue + + element = item.value if isinstance(item, ast.DictComp) else item.elt + new_result = copy.deepcopy(item) + + if isinstance(element, Placeholder): + loop_iter = ensure_ast(iterator) + elif isinstance(element, ast.IfExp) and any( + isinstance(x, Placeholder) for x in ast.walk(element) + ): + # A conditional expression was being built up in the element slot, + # but it turned out to be this loop's iterable, as in + # `for y in (a if c else b)`. Move it back out and plug the value + # this path produced into the arm still awaiting one. + if isinstance(new_result, ast.DictComp): + new_result.key, new_result.value = Placeholder(), Placeholder() + else: + new_result.elt = Placeholder() + + plugged = ReplacePlaceholder(ensure_ast(iterator)).visit( + copy.deepcopy(element) + ) + assert isinstance(plugged, ast.expr) + loop_iter = plugged + else: + continue + + # The loop target is not named until the STORE_* that follows. + loop_info = ast.comprehension( + target=TargetHole(), iter=loop_iter, ifs=[], is_async=0 + ) + new_result.generators.append(loop_info) + new_stack = ( + state.stack[:pos] + + [new_result] + + state.stack[pos + 1 :] + + [loop_info.target] + ) + return replace(state, stack=new_stack) + + raise TypeError("FOR_ITER did not find partial comprehension on stack") + + +@register_handler("GET_ITER", version=PythonVersion.PY_312) +@register_handler("GET_ITER", version=PythonVersion.PY_313) +@register_handler("GET_ITER", version=PythonVersion.PY_314) +def handle_get_iter( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # GET_ITER converts the top stack item to an iterator + # For AST reconstruction, we typically don't need to change anything + # since the iterator will be used directly in the comprehension + return state + + +@register_handler("END_FOR", version=PythonVersion.PY_312) +def handle_end_for_312( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # END_FOR marks the end of a for loop, followed by POP_TOP (in 3.12) + new_stack = state.stack[:-2] + return replace(state, stack=new_stack) + + +@register_handler("END_FOR", version=PythonVersion.PY_313) +@register_handler("END_FOR", version=PythonVersion.PY_314) +def handle_end_for( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # END_FOR marks the end of a for loop - no action needed for AST reconstruction + new_stack = state.stack[:-1] + return replace(state, stack=new_stack) + + +@register_handler("RERAISE", version=PythonVersion.PY_312) +@register_handler("RERAISE", version=PythonVersion.PY_313) +@register_handler("RERAISE", version=PythonVersion.PY_314) +def handle_reraise( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # RERAISE re-raises an exception - generally ignore for AST reconstruction + return state + + +# ============================================================================ +# VARIABLE OPERATIONS HANDLERS +# ============================================================================ + + +def _literal_elements(value: ast.expr, count: int | None = None) -> list[ast.expr]: + """The elements of a literal sequence, for destructuring a known value.""" + assert isinstance(value, ast.Tuple | ast.List), ( + f"Cannot unpack {type(value).__name__}; expected a literal sequence" + ) + assert count is None or len(value.elts) == count, ( + f"Expected {count} values to unpack, got {len(value.elts)}" + ) + return [ensure_ast(element) for element in value.elts] + + +def _bind_local( + state: ReconstructionState, var_name: str, value: ast.expr +) -> ReconstructionState: + """Record that a local now stands for ``value``, popping it off the stack. + + Reached when a store is not filling in a loop target: 3.14 unrolls a + single-iteration loop over a literal, assigning its targets outright. The + loop is gone from the bytecode, so the names it bound are not in scope in + the reconstruction and their uses are substituted instead. + """ + bindings = {**state.bindings, var_name: ensure_ast(value)} + return replace(state, stack=state.stack[:-1], bindings=bindings) + + +def _read_local(state: ReconstructionState, var_name: str) -> ast.expr: + """The expression a local name stands for at this point on this path.""" + if var_name == DummyIterName().id: + return DummyIterName() + elif var_name in state.bindings: + # Bound to a known expression rather than by a loop, so the name itself + # is not in scope in the reconstruction; use what it was bound to. + return copy.deepcopy(state.bindings[var_name]) + else: + return ast.Name(id=var_name, ctx=ast.Load()) + + +@register_handler("LOAD_FAST", version=PythonVersion.PY_312) +@register_handler("LOAD_FAST", version=PythonVersion.PY_313) +@register_handler("LOAD_FAST", version=PythonVersion.PY_314) +@register_handler("LOAD_FAST_CHECK", version=PythonVersion.PY_312) +@register_handler("LOAD_FAST_CHECK", version=PythonVersion.PY_313) +@register_handler("LOAD_FAST_CHECK", version=PythonVersion.PY_314) +def handle_load_fast( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # LOAD_FAST_CHECK differs only in raising when the local is unbound, which + # says nothing about the expression being reconstructed. + return replace(state, stack=state.stack + [_read_local(state, instr.argval)]) + + +@register_handler("LOAD_DEREF", version=PythonVersion.PY_312) +@register_handler("LOAD_DEREF", version=PythonVersion.PY_313) +@register_handler("LOAD_DEREF", version=PythonVersion.PY_314) +def handle_load_deref( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # LOAD_DEREF loads a value out of a cell. When the cell is one this code + # object *creates* (a comprehension target captured by a nested lambda, say) + # the name is bound inside the reconstructed tree too, so it can stand. A + # free variable is different: its cell belongs to an enclosing scope that + # the reconstructed tree does not reproduce, so a bare name would silently + # become a global lookup. Write the captured value in instead. + var_name = instr.argval + if var_name in state.code.co_freevars and var_name in state.freevars: + loaded: ast.expr = copy.deepcopy(state.freevars[var_name]) + else: + loaded = ast.Name(id=var_name, ctx=ast.Load()) + return replace(state, stack=state.stack + [loaded]) + + +@register_handler("LOAD_CLOSURE", version=PythonVersion.PY_312) +@register_handler("LOAD_CLOSURE", version=PythonVersion.PY_313) +@register_handler("LOAD_CLOSURE", version=PythonVersion.PY_314) +def handle_load_closure( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # LOAD_CLOSURE loads a closure variable + var_name = instr.argval + new_stack = state.stack + [ast.Name(id=var_name, ctx=ast.Load())] + return replace(state, stack=new_stack) + + +@register_handler("LOAD_CONST", version=PythonVersion.PY_312) +@register_handler("LOAD_CONST", version=PythonVersion.PY_313) +@register_handler("LOAD_CONST", version=PythonVersion.PY_314) +def handle_load_const( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + const_value = instr.argval + # A nested lambda or comprehension arrives as a code object, and its own + # free variables reach back through this scope to the same cells, so the + # captured values have to travel with it. + loaded = ( + _reconstruct_code(const_value, state.freevars) + if isinstance(const_value, types.CodeType) + else ensure_ast(const_value) + ) + return replace(state, stack=state.stack + [loaded]) + + +@register_handler("LOAD_GLOBAL", version=PythonVersion.PY_312) +@register_handler("LOAD_GLOBAL", version=PythonVersion.PY_313) +@register_handler("LOAD_GLOBAL", version=PythonVersion.PY_314) +def handle_load_global( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + global_name = instr.argval + + if instr.argrepr.endswith(" + NULL"): + new_stack = state.stack + [ast.Name(id=global_name, ctx=ast.Load()), Null()] + elif instr.argrepr.startswith("NULL + "): + new_stack = state.stack + [Null(), ast.Name(id=global_name, ctx=ast.Load())] + else: + new_stack = state.stack + [ast.Name(id=global_name, ctx=ast.Load())] + return replace(state, stack=new_stack) + + +@register_handler("LOAD_NAME", version=PythonVersion.PY_312) +@register_handler("LOAD_NAME", version=PythonVersion.PY_313) +@register_handler("LOAD_NAME", version=PythonVersion.PY_314) +def handle_load_name( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # LOAD_NAME is similar to LOAD_GLOBAL but for names in the global namespace + name = instr.argval + new_stack = state.stack + [ast.Name(id=name, ctx=ast.Load())] + return replace(state, stack=new_stack) + + +def _is_assignment_expression(state: ReconstructionState) -> bool: + """Is the top of the stack a COPY of the value beneath it? + + That duplication is how an assignment expression keeps its value after + binding it: `COPY 1` then a `STORE_*`. `handle_copy` pushes the very same + node, so identity is what distinguishes it from two equal-looking values. + """ + return len(state.stack) >= 2 and state.stack[-1] is state.stack[-2] + + +def _handle_assignment_expression( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + """Rebuild `(name := value)` from the COPY/STORE pair that implements it.""" + target = ast.Name(id=instr.argval, ctx=ast.Store()) + named = ast.NamedExpr(target=target, value=ensure_ast(state.stack[-1])) + return replace(state, stack=state.stack[:-2] + [named]) + + +@register_handler("STORE_GLOBAL", version=PythonVersion.PY_312) +@register_handler("STORE_GLOBAL", version=PythonVersion.PY_313) +@register_handler("STORE_GLOBAL", version=PythonVersion.PY_314) +def handle_store_global( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # A comprehension has no globals of its own, so the only way it stores one + # is an assignment expression, which binds in the enclosing scope. + assert _is_assignment_expression(state), ( + "STORE_GLOBAL outside an assignment expression" + ) + return _handle_assignment_expression(state, instr) + + +@register_handler("STORE_DEREF", version=PythonVersion.PY_312) +@register_handler("STORE_DEREF", version=PythonVersion.PY_313) +@register_handler("STORE_DEREF", version=PythonVersion.PY_314) +def handle_store_deref( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # STORE_DEREF stores into a closure variable: either an assignment + # expression binding in an enclosing function, or a loop target that an + # inner comprehension captures. + if _is_assignment_expression(state): + return _handle_assignment_expression(state, instr) + return handle_store_fast(state, instr) + + +@register_handler("STORE_FAST", version=PythonVersion.PY_312) +@register_handler("STORE_FAST", version=PythonVersion.PY_313) +@register_handler("STORE_FAST", version=PythonVersion.PY_314) +def handle_store_fast( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + if _is_assignment_expression(state): + # An assignment expression whose target is local to this code object, + # as in a comprehension inlined into the lambda that binds the name. + return _handle_assignment_expression(state, instr) + + if isinstance(state.stack[-1], ast.Name) and state.stack[-1].id == instr.argval: + # If the variable is already on the stack, we can skip adding it again + # This is common in nested comprehensions where the same variable is reused + return replace(state, stack=state.stack[:-1]) + + if not isinstance(state.stack[-1], TargetHole): + return _bind_local(state, instr.argval, state.stack[-1]) + + new_stack = _bind_target_hole( + state.stack, state.stack[-1], ast.Name(id=instr.argval, ctx=ast.Store()) + ) + return replace(state, stack=new_stack[:-1]) + + +@register_handler("STORE_FAST_LOAD_FAST", version=PythonVersion.PY_313) +@register_handler("STORE_FAST_LOAD_FAST", version=PythonVersion.PY_314) +def handle_store_fast_load_fast( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # STORE_FAST_LOAD_FAST stores and then loads the same variable (optimization) + # The instruction has two names: store_name and load_name + # In Python 3.13, this is often used for loop variables + + # In Python 3.13, the instruction argument contains both names + # argval should be a tuple (store_name, load_name) + assert isinstance(instr.argval, tuple) + store_name, load_name = instr.argval + + if _is_assignment_expression(state): + # `(name := value)` whose result is read straight back, as in + # `(z := w) + z`. The duplicate becomes the assignment expression and + # the reload becomes a plain reference to the name just bound. + named = ast.NamedExpr( + target=ast.Name(id=store_name, ctx=ast.Store()), + value=ensure_ast(state.stack[-1]), + ) + reload = ast.Name(id=load_name, ctx=ast.Load()) + return replace(state, stack=state.stack[:-2] + [named, reload]) + + if not isinstance(state.stack[-1], TargetHole): + # A plain assignment followed by a load, as 3.14 emits when it unrolls + # a single-iteration loop over a literal. + bound = _bind_local(state, store_name, state.stack[-1]) + return replace(bound, stack=bound.stack + [_read_local(bound, load_name)]) + + new_stack = _bind_target_hole( + state.stack, state.stack[-1], ast.Name(id=store_name, ctx=ast.Store()) + ) + new_var = ast.Name(id=load_name, ctx=ast.Load()) + return replace(state, stack=new_stack[:-1] + [new_var]) + + +@register_handler("STORE_FAST_STORE_FAST", version=PythonVersion.PY_313) +@register_handler("STORE_FAST_STORE_FAST", version=PythonVersion.PY_314) +def handle_store_fast_store_fast( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # STORE_FAST_STORE_FAST stores STACK[-1] into the first named variable and + # STACK[-2] into the second. It is emitted for unpacking targets, so both + # values are loop-target holes belonging to the same comprehension. + assert isinstance(instr.argval, tuple) + first_name, second_name = instr.argval + + if not isinstance(state.stack[-1], TargetHole): + # Not loop targets: a pair of plain assignments, as 3.14 emits when it + # unrolls a single-iteration loop over a literal. + bound = _bind_local(state, first_name, state.stack[-1]) + return _bind_local(bound, second_name, bound.stack[-1]) + + new_stack = _bind_target_hole( + state.stack, state.stack[-1], ast.Name(id=first_name, ctx=ast.Store()) + ) + new_stack = _bind_target_hole( + new_stack, new_stack[-2], ast.Name(id=second_name, ctx=ast.Store()) + ) + return replace(state, stack=new_stack[:-2]) + + +@register_handler("LOAD_FAST_AND_CLEAR", version=PythonVersion.PY_312) +@register_handler("LOAD_FAST_AND_CLEAR", version=PythonVersion.PY_313) +@register_handler("LOAD_FAST_AND_CLEAR", version=PythonVersion.PY_314) +def handle_load_fast_and_clear( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # LOAD_FAST_AND_CLEAR pushes a local variable onto the stack and clears it + # For AST reconstruction, we treat this the same as LOAD_FAST + return replace(state, stack=state.stack + [_read_local(state, instr.argval)]) + + +@register_handler("LOAD_FAST_LOAD_FAST", version=PythonVersion.PY_313) +@register_handler("LOAD_FAST_LOAD_FAST", version=PythonVersion.PY_314) +def handle_load_fast_load_fast( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # LOAD_FAST_LOAD_FAST loads two variables (optimization in Python 3.13) + # The instruction argument contains both variable names + if isinstance(instr.argval, tuple): + var1, var2 = instr.argval + else: + # Fallback: assume both names are the same + var1 = var2 = instr.argval + + new_stack = state.stack + [_read_local(state, var1), _read_local(state, var2)] + + return replace(state, stack=new_stack) + + +@register_handler("MAKE_CELL", version=PythonVersion.PY_312) +@register_handler("MAKE_CELL", version=PythonVersion.PY_313) +@register_handler("MAKE_CELL", version=PythonVersion.PY_314) +def handle_make_cell( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # MAKE_CELL creates a new cell in slot i for closure variables + # This is used when variables from outer scopes are captured by inner scopes + # For AST reconstruction purposes, this is just a variable scoping mechanism + # that we can ignore since the AST doesn't track low-level closure details + return state + + +@register_handler("COPY_FREE_VARS", version=PythonVersion.PY_312) +@register_handler("COPY_FREE_VARS", version=PythonVersion.PY_313) +@register_handler("COPY_FREE_VARS", version=PythonVersion.PY_314) +def handle_copy_free_vars( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # COPY_FREE_VARS copies n free (closure) variables from the closure into the frame + # This removes the need for special code on the caller's side when calling closures + # For AST reconstruction purposes, this is just a variable scoping mechanism + # that we can ignore since the AST doesn't track runtime variable management + return state + + +# ============================================================================ +# STACK MANAGEMENT HANDLERS +# ============================================================================ + + +@register_handler("POP_TOP", version=PythonVersion.PY_312) +@register_handler("POP_TOP", version=PythonVersion.PY_313) +@register_handler("POP_TOP", version=PythonVersion.PY_314) +def handle_pop_top( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # POP_TOP removes the top item from the stack + # In generators, often used after YIELD_VALUE + # Also used to clean up the duplicated middle value in failed chained comparisons + new_stack = state.stack[:-1] + return replace(state, stack=new_stack) + + +# Python 3.13 replacement for stack manipulation +@register_handler("SWAP", version=PythonVersion.PY_312) +@register_handler("SWAP", version=PythonVersion.PY_313) +@register_handler("SWAP", version=PythonVersion.PY_314) +def handle_swap( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # SWAP exchanges the top two stack items (replaces ROT_TWO in many cases) + assert instr.arg is not None + depth = instr.arg + stack_size = len(state.stack) + + if depth > stack_size: + # Not enough items on stack - this might be a pattern where some items were optimized away + # For AST reconstruction, we can often ignore certain stack manipulations + return state + + # For other depths, swap TOS with the item at specified depth + assert depth <= stack_size, f"SWAP depth {depth} exceeds stack size {stack_size}" + idx = stack_size - depth + new_stack = state.stack.copy() + new_stack[-1], new_stack[idx] = new_stack[idx], new_stack[-1] + return replace(state, stack=new_stack) + + +@register_handler("COPY", version=PythonVersion.PY_312) +@register_handler("COPY", version=PythonVersion.PY_313) +@register_handler("COPY", version=PythonVersion.PY_314) +def handle_copy( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # COPY duplicates the item at the specified depth + assert instr.arg is not None + depth = instr.arg + stack_size = len(state.stack) + if depth > stack_size: + raise ValueError(f"COPY depth {depth} exceeds stack size {stack_size}") + idx = stack_size - depth + copied_item = state.stack[idx] + new_stack = state.stack + [copied_item] + return replace(state, stack=new_stack) + + +@register_handler("PUSH_NULL", version=PythonVersion.PY_312) +@register_handler("PUSH_NULL", version=PythonVersion.PY_313) +@register_handler("PUSH_NULL", version=PythonVersion.PY_314) +def handle_push_null( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + return replace(state, stack=state.stack + [Null()]) + + +# ============================================================================ +# BINARY ARITHMETIC/LOGIC OPERATION HANDLERS +# ============================================================================ + + +def handle_binop( + op: ast.operator, state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + right = ensure_ast(state.stack[-1]) + left = ensure_ast(state.stack[-2]) + new_stack = state.stack[:-2] + [ast.BinOp(left=left, op=op, right=right)] + return replace(state, stack=new_stack) + + +# Python 3.12+ BINARY_OP handler +@register_handler("BINARY_OP", version=PythonVersion.PY_312) +@register_handler("BINARY_OP", version=PythonVersion.PY_313) +def handle_binary_op( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # BINARY_OP in Python 3.12+ consolidates all binary operations + # The operation type is determined by the instruction argument + assert instr.arg is not None + + # Map argument values to AST operators based on Python 3.12+ implementation + op_map: collections.abc.Mapping[int, ast.operator] = { + 0: ast.Add(), # + + 1: ast.BitAnd(), # & + 2: ast.FloorDiv(), # // + 3: ast.LShift(), # << + 4: ast.MatMult(), # @ + 5: ast.Mult(), # * + 6: ast.Mod(), # % + 7: ast.BitOr(), # | + 8: ast.Pow(), # ** + 9: ast.RShift(), # >> + 10: ast.Sub(), # - + 11: ast.Div(), # / + 12: ast.BitXor(), # ^ + } + + op = op_map.get(instr.arg) + if op is None: + raise TypeError(f"Unknown binary operation: {instr.arg}") + + return handle_binop(op, state, instr) + + +# 3.14 folded subscripting into BINARY_OP; `dis._nb_ops` names the oparg +# NB_SUBSCR. Looked up rather than hard-coded, since it sits past the in-place +# operators and so moves whenever one is added. +_NB_OPS: list[tuple[str, str]] = getattr(dis, "_nb_ops", []) +NB_SUBSCR: int | None = next( + (i for i, (name, _) in enumerate(_NB_OPS) if name == "NB_SUBSCR"), None +) + + +@register_handler("BINARY_OP", version=PythonVersion.PY_314) +def handle_binary_op_314( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # As in 3.13, except that BINARY_OP now also implements `a[b]`, which used + # to be its own BINARY_SUBSCR opcode. + if instr.arg is not None and instr.arg == NB_SUBSCR: + return handle_binary_subscr(state, instr) + return handle_binary_op(state, instr) + + +@register_handler("LOAD_SMALL_INT", version=PythonVersion.PY_314) +def handle_load_small_int( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # LOAD_SMALL_INT pushes an int in range(256) held in the oparg itself, + # rather than going through co_consts. + assert isinstance(instr.argval, int) + return replace(state, stack=state.stack + [ensure_ast(instr.argval)]) + + +@register_handler("LOAD_FAST_BORROW", version=PythonVersion.PY_314) +def handle_load_fast_borrow( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # A borrowed reference differs only in ownership, which the AST does not + # model, so this is LOAD_FAST as far as reconstruction is concerned. + return handle_load_fast(state, instr) + + +@register_handler("LOAD_FAST_BORROW_LOAD_FAST_BORROW", version=PythonVersion.PY_314) +def handle_load_fast_borrow_load_fast_borrow( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + return handle_load_fast_load_fast(state, instr) + + +@register_handler("LOAD_COMMON_CONSTANT", version=PythonVersion.PY_314) +def handle_load_common_constant( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # Pushes one of a small hardcoded set of constants. In a comprehension this + # only shows up in the guard of an inlined builtin; see CommonConstant. + name = getattr(instr.argval, "__name__", str(instr.argval)) + return replace(state, stack=state.stack + [CommonConstant(id=name)]) + + +@register_handler("NOT_TAKEN", version=PythonVersion.PY_314) +def handle_not_taken( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # A no-op marking the not-taken edge of a branch for sys.monitoring. + return state + + +@register_handler("POP_ITER", version=PythonVersion.PY_314) +def handle_pop_iter( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # POP_ITER discards the exhausted iterator that FOR_ITER left behind. In + # 3.13 the same cleanup was spelled END_FOR followed by POP_TOP. + return replace(state, stack=state.stack[:-1]) + + +# ============================================================================ +# UNARY OPERATION HANDLERS +# ============================================================================ + + +def handle_unary_op( + op: ast.unaryop, state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + operand = ensure_ast(state.stack[-1]) + new_stack = state.stack[:-1] + [ast.UnaryOp(op=op, operand=operand)] + return replace(state, stack=new_stack) + + +UNARY_OPS: dict[str, ast.unaryop] = { + "UNARY_NEGATIVE": ast.USub(), + "UNARY_INVERT": ast.Invert(), + "UNARY_NOT": ast.Not(), +} + +# These three behave identically in every dialect this module supports; 3.13's +# "requires an exact bool operand" note on UNARY_NOT constrains the operand, not +# the reconstruction. +for _opname, _op in UNARY_OPS.items(): + for _version in ( + PythonVersion.PY_312, + PythonVersion.PY_313, + PythonVersion.PY_314, + ): + register_handler( + _opname, functools.partial(handle_unary_op, _op), version=_version + ) + + +@register_handler("CONVERT_VALUE", version=PythonVersion.PY_313) +@register_handler("CONVERT_VALUE", version=PythonVersion.PY_314) +def handle_convert_value( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # CONVERT_VALUE applies a conversion to the value on top of stack + # Used for f-string conversions like !r, !s, !a + # The conversion type is stored in instr.arg: + # 0 = None, 1 = str (!s), 2 = repr (!r), 3 = ascii (!a) + assert len(state.stack) > 0, "CONVERT_VALUE requires a value on stack" + assert instr.arg is not None, "CONVERT_VALUE requires conversion type" + + # Wrap the value with conversion information + value = state.stack[-1] + converted = ConvertedValue(value, instr.arg) + new_stack = state.stack[:-1] + [converted] + + return replace(state, stack=new_stack) + + +@register_handler("CALL_INTRINSIC_1", version=PythonVersion.PY_312) +@register_handler("CALL_INTRINSIC_1", version=PythonVersion.PY_313) +@register_handler("CALL_INTRINSIC_1", version=PythonVersion.PY_314) +def handle_call_intrinsic_1( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # CALL_INTRINSIC_1 calls an intrinsic function with one argument + if instr.argrepr == "INTRINSIC_LIST_TO_TUPLE": + assert isinstance(state.stack[-1], ast.List), ( + "Expected a list for LIST_TO_TUPLE" + ) + tuple_node = ast.Tuple(elts=state.stack[-1].elts, ctx=ast.Load()) + return replace(state, stack=state.stack[:-1] + [tuple_node]) + elif instr.argrepr == "INTRINSIC_UNARY_POSITIVE": + assert len(state.stack) > 0 + new_val = ast.UnaryOp(op=ast.UAdd(), operand=state.stack[-1]) + return replace(state, stack=state.stack[:-1] + [new_val]) + elif instr.argrepr == "INTRINSIC_STOPITERATION_ERROR": + return state + else: + raise TypeError(f"Unsupported generator intrinsic operation: {instr.argrepr}") + + +@register_handler("TO_BOOL", version=PythonVersion.PY_313) +@register_handler("TO_BOOL", version=PythonVersion.PY_314) +def handle_to_bool( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # TO_BOOL converts the top stack item to a boolean + # For AST reconstruction, we typically don't need an explicit bool() call + # since the boolean context is usually handled by the conditional jump that follows + # However, for some cases we might need to preserve the explicit conversion + + # For now, leave the value as-is since the jump instruction will handle the boolean logic + return state + + +# ============================================================================ +# COMPARISON OPERATION HANDLERS +# ============================================================================ + +CMP_OPMAP: dict[str, ast.cmpop] = { + "<": ast.Lt(), + "<=": ast.LtE(), + ">": ast.Gt(), + ">=": ast.GtE(), + "==": ast.Eq(), + "!=": ast.NotEq(), +} + + +@register_handler("COMPARE_OP", version=PythonVersion.PY_312) +@register_handler("COMPARE_OP", version=PythonVersion.PY_313) +@register_handler("COMPARE_OP", version=PythonVersion.PY_314) +def handle_compare_op( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + assert instr.arg is not None and instr.argval in dis.cmp_op, ( + f"Unsupported comparison operation: {instr.argval}" + ) + + right = ensure_ast(state.stack[-1]) + left = ensure_ast(state.stack[-2]) + + # Map comparison operation codes to AST operators + op_name = instr.argval + compare_node = ast.Compare(left=left, ops=[CMP_OPMAP[op_name]], comparators=[right]) + new_stack = state.stack[:-2] + [compare_node] + return replace(state, stack=new_stack) + + +@register_handler("CONTAINS_OP", version=PythonVersion.PY_312) +@register_handler("CONTAINS_OP", version=PythonVersion.PY_313) +@register_handler("CONTAINS_OP", version=PythonVersion.PY_314) +def handle_contains_op( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + right = ensure_ast(state.stack[-1]) # Container + left = ensure_ast(state.stack[-2]) # Item to check + + # instr.arg determines if it's 'in' (0) or 'not in' (1) + op = ast.NotIn() if instr.arg else ast.In() + + compare_node = ast.Compare(left=left, ops=[op], comparators=[right]) + new_stack = state.stack[:-2] + [compare_node] + return replace(state, stack=new_stack) + + +@register_handler("IS_OP", version=PythonVersion.PY_312) +@register_handler("IS_OP", version=PythonVersion.PY_313) +@register_handler("IS_OP", version=PythonVersion.PY_314) +def handle_is_op( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + right = ensure_ast(state.stack[-1]) + left = ensure_ast(state.stack[-2]) + + # instr.arg determines if it's 'is' (0) or 'is not' (1) + op = ast.IsNot() if instr.arg else ast.Is() + + compare_node = ast.Compare(left=left, ops=[op], comparators=[right]) + new_stack = state.stack[:-2] + [compare_node] + return replace(state, stack=new_stack) + + +# ============================================================================ +# FUNCTION CALL HANDLERS +# ============================================================================ + + +@register_handler("KW_NAMES", version=PythonVersion.PY_312) +def handle_kw_names( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # KW_NAMES names the trailing arguments of the CALL that follows it. + # Python 3.13 replaced this pair with a single CALL_KW instruction. + assert isinstance(instr.argval, tuple), "KW_NAMES requires a tuple of names" + assert all(isinstance(name, str) for name in instr.argval) + assert state.kw_names is None, "KW_NAMES must be consumed by the following CALL" + return replace(state, kw_names=instr.argval) + + +@register_handler("CALL", version=PythonVersion.PY_312) +def handle_call_312( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # CALL in Python 3.12 handles both function and method calls + # Stack layout: [..., callable or self, callable or NULL] + assert instr.arg is not None + arg_count: int = instr.arg + + # Check if this is a method call (no NULL on top) + if isinstance(state.stack[-arg_count - 2], Null): + # Regular function call: [..., NULL, callable, *args] + func = ensure_ast(state.stack[-arg_count - 1]) + args = ( + [ensure_ast(arg) for arg in state.stack[-arg_count:]] + if arg_count > 0 + else [] + ) + new_stack = state.stack[: -arg_count - 2] + else: + # Method call: [..., callable, self, *args] + func = ensure_ast(state.stack[-arg_count - 2]) + self_arg = ensure_ast(state.stack[-arg_count - 1]) + remaining_args = ( + [ensure_ast(arg) for arg in state.stack[-arg_count:]] + if arg_count > 0 + else [] + ) + args = [self_arg] + remaining_args + new_stack = state.stack[: -arg_count - 2] + + # A preceding KW_NAMES names the trailing `len(kw_names)` positional slots. + keywords: list[ast.keyword] = [] + if state.kw_names is not None: + assert 0 < len(state.kw_names) <= arg_count + keywords = [ + ast.keyword(arg=name, value=value) + for name, value in zip(state.kw_names, args[-len(state.kw_names) :]) + ] + args = args[: -len(state.kw_names)] + + if isinstance(func, CompLambda): + assert len(args) == 1 and not keywords + return replace(state, stack=new_stack + [func.inline(args[0])], kw_names=None) + else: + # Create function call AST + call_node = ast.Call(func=func, args=args, keywords=keywords) + new_stack = new_stack + [call_node] + return replace(state, stack=new_stack, kw_names=None) + + +@register_handler("CALL", version=PythonVersion.PY_313) +@register_handler("CALL", version=PythonVersion.PY_314) +def handle_call( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # CALL pops function and arguments from stack (replaces CALL_FUNCTION in Python 3.13) + assert instr.arg is not None + arg_count: int = instr.arg + + func = ensure_ast(state.stack[-arg_count - 2]) + + # Pop arguments and function + args = ( + [ensure_ast(arg) for arg in state.stack[-arg_count:]] if arg_count > 0 else [] + ) + if not isinstance(state.stack[-arg_count - 1], Null): + args = [ensure_ast(state.stack[-arg_count - 1])] + args + + new_stack = state.stack[: -arg_count - 2] + if isinstance(func, CompLambda): + assert len(args) == 1 + return replace(state, stack=new_stack + [func.inline(args[0])]) + else: + # Create function call AST + call_node = ast.Call(func=func, args=args, keywords=[]) + new_stack = new_stack + [call_node] + return replace(state, stack=new_stack) + + +@register_handler("CALL_KW", version=PythonVersion.PY_313) +@register_handler("CALL_KW", version=PythonVersion.PY_314) +def handle_call_kw( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # CALL_KW pops function, arguments, and keyword names from stack + assert instr.arg is not None + arg_count: int = instr.arg + assert arg_count > 0, "CALL_KW requires at least one argument" + + func = ensure_ast(state.stack[-arg_count - 3]) + assert not isinstance(func, CompLambda | Null) + + kw_names = state.stack[-1] + assert isinstance(kw_names, ast.Tuple), "Expected a tuple of keyword names" + assert len(kw_names.elts) > 0, "Expected at least one keyword name" + + # Pop arguments, function, and keyword names + keywords = [] + for i, kw in enumerate(reversed(kw_names.elts)): + assert isinstance(kw, ast.Constant) and isinstance(kw.value, str) + keywords += [ast.keyword(arg=kw.value, value=ensure_ast(state.stack[-2 - i]))] + keywords.reverse() + + args = [ensure_ast(a) for a in state.stack[-arg_count - 1 : -len(keywords) - 1]] + if not isinstance(state.stack[-arg_count - 2], Null): + args = [ensure_ast(state.stack[-arg_count - 2])] + args + + # Create function call AST + call_node = ast.Call(func=func, args=args, keywords=keywords) + new_stack = state.stack[: -arg_count - 3] + [call_node] + return replace(state, stack=new_stack) + + +# Flags shared by MAKE_FUNCTION (3.12) and SET_FUNCTION_ATTRIBUTE (3.13) +MAKE_FUNCTION_DEFAULTS = 0x01 +MAKE_FUNCTION_KWDEFAULTS = 0x02 +MAKE_FUNCTION_ANNOTATIONS = 0x04 +MAKE_FUNCTION_CLOSURE = 0x08 +MAKE_FUNCTION_ANNOTATE = 0x10 # added in 3.14 +MAKE_FUNCTION_FLAGS = ( + MAKE_FUNCTION_DEFAULTS, + MAKE_FUNCTION_KWDEFAULTS, + MAKE_FUNCTION_ANNOTATIONS, + MAKE_FUNCTION_CLOSURE, +) + + +def _apply_function_attribute( + func: ast.Lambda | CompLambda, flag: int, value: ast.expr +) -> ast.Lambda | CompLambda: + """Attach one function attribute to a reconstructed lambda.""" + if flag == MAKE_FUNCTION_CLOSURE: + # The body has already resolved each free variable: to a captured value + # if the cell came from outside the comprehension, and otherwise to the + # name the reconstructed tree binds it under. See `handle_load_deref`. + return func + if flag == MAKE_FUNCTION_ANNOTATE: + # A lambda has no annotations, and the AST does not carry the lazy + # annotate function 3.14 attaches to annotated functions. + return func + + assert isinstance(func, ast.Lambda) and not isinstance(func, CompLambda), ( + "Only lambdas carry defaults; comprehensions take exactly one argument" + ) + + if flag == MAKE_FUNCTION_DEFAULTS: + # A tuple of defaults for the *trailing* positional parameters. + assert isinstance(value, ast.Tuple), "Expected a tuple of default values" + func.args.defaults = list(value.elts) + elif flag == MAKE_FUNCTION_KWDEFAULTS: + # A dict mapping keyword-only parameter names to their defaults. + assert isinstance(value, ast.Dict), "Expected a dict of keyword defaults" + by_name = { + key.value: val + for key, val in zip(value.keys, value.values) + if isinstance(key, ast.Constant) + } + func.args.kw_defaults = [by_name.get(a.arg) for a in func.args.kwonlyargs] + else: + raise NotImplementedError("Function annotations are not supported") + + return func + + +def _split_callable( + first: ast.expr, second: ast.expr +) -> tuple[ast.expr, ast.expr | None]: + """Separate the callable from the NULL-or-self slot beside it. + + Which of the two comes first varies: LOAD_GLOBAL and LOAD_ATTR report the + order in their argrepr, and it differs between 3.13 and 3.14. + """ + if isinstance(first, Null): + return second, None + elif isinstance(second, Null): + return first, None + else: + return first, second + + +def _build_variadic_call( + func: ast.expr, + self_arg: ast.expr | None, + positional: ast.expr, + keyword_mapping: ast.expr | None, +) -> ast.Call: + """Assemble `func(*positional, **keyword_mapping)`. + + CALL_FUNCTION_EX receives its arguments already collected into a sequence + and a mapping, with the original mix of plain and starred arguments no + longer distinguishable. Spelling every argument as unpacked reproduces the + call exactly, even where the source did not use `*` for all of them. + """ + args: list[ast.expr] = [] if self_arg is None else [ensure_ast(self_arg)] + args.append(ast.Starred(value=ensure_ast(positional), ctx=ast.Load())) + keywords = ( + [] + if keyword_mapping is None + else [ast.keyword(arg=None, value=ensure_ast(keyword_mapping))] + ) + return ast.Call(func=ensure_ast(func), args=args, keywords=keywords) + + +@register_handler("CALL_FUNCTION_EX", version=PythonVersion.PY_312) +@register_handler("CALL_FUNCTION_EX", version=PythonVersion.PY_313) +def handle_call_function_ex( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # Stack: callable and NULL-or-self, the positional sequence, and -- only + # when the low bit of the oparg is set -- the keyword mapping. + size = 4 if instr.arg else 3 + keyword_mapping = state.stack[-1] if instr.arg else None + positional = state.stack[-2] if instr.arg else state.stack[-1] + func, self_arg = _split_callable(state.stack[-size], state.stack[-size + 1]) + + call = _build_variadic_call(func, self_arg, positional, keyword_mapping) + return replace(state, stack=state.stack[:-size] + [call]) + + +@register_handler("CALL_FUNCTION_EX", version=PythonVersion.PY_314) +def handle_call_function_ex_314( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # 3.14 always reserves the keyword-mapping slot, pushing NULL into it when + # the call has no `**` argument, so the layout is a fixed four slots. + keyword_mapping = None if isinstance(state.stack[-1], Null) else state.stack[-1] + func, self_arg = _split_callable(state.stack[-4], state.stack[-3]) + + call = _build_variadic_call(func, self_arg, state.stack[-2], keyword_mapping) + return replace(state, stack=state.stack[:-4] + [call]) + + +@register_handler("MAKE_FUNCTION", version=PythonVersion.PY_312) +def handle_make_function_312( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # MAKE_FUNCTION in Python 3.12 uses flags to determine stack consumption. + # Unlike 3.10 there is no qualified name on the stack, and unlike 3.13 the + # extra attributes travel with this instruction rather than with a following + # SET_FUNCTION_ATTRIBUTE. They are pushed in ascending flag order, below the + # code object. + assert instr.arg is not None + assert isinstance(state.stack[-1], ast.Lambda | CompLambda), ( + "Expected a function object (Lambda or CompLambda) on the stack." + ) + + set_flags = [flag for flag in MAKE_FUNCTION_FLAGS if instr.arg & flag] + attributes = state.stack[-1 - len(set_flags) : -1] + + func = copy.deepcopy(state.stack[-1]) + for flag, value in zip(set_flags, attributes): + func = _apply_function_attribute(func, flag, value) + + new_stack = state.stack[: -1 - len(set_flags)] + [func] + return replace(state, stack=new_stack) + + +# Python 3.13 version +@register_handler("MAKE_FUNCTION", version=PythonVersion.PY_313) +@register_handler("MAKE_FUNCTION", version=PythonVersion.PY_314) +def handle_make_function( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # MAKE_FUNCTION in Python 3.13 is simplified: it only takes a code object from the stack + # and creates a function from it. No flags, no extra attributes on the stack. + # All extra attributes are handled by separate SET_FUNCTION_ATTRIBUTE instructions. + + # Pop the function object from the stack (it's the only thing expected) + # Conversion from CodeType to ast.Lambda should have happened already + assert isinstance(state.stack[-1], ast.Lambda | CompLambda), ( + "Expected a function object (Lambda or CompLambda) on the stack." + ) + return state + + +@register_handler("SET_FUNCTION_ATTRIBUTE", version=PythonVersion.PY_313) +@register_handler("SET_FUNCTION_ATTRIBUTE", version=PythonVersion.PY_314) +def handle_set_function_attribute( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # SET_FUNCTION_ATTRIBUTE sets one attribute on a function object. Python + # 3.13 uses it in place of the MAKE_FUNCTION flags; the stack holds the + # attribute value below the function, and only the function is left behind. + assert instr.arg is not None + assert isinstance(state.stack[-1], ast.Lambda | CompLambda), ( + "Expected a function object (Lambda or CompLambda) on the stack." + ) + + func = _apply_function_attribute( + copy.deepcopy(state.stack[-1]), instr.arg, state.stack[-2] + ) + return replace(state, stack=state.stack[:-2] + [func]) + + +# ============================================================================ +# OBJECT ACCESS HANDLERS +# ============================================================================ + + +@register_handler("LOAD_ATTR", version=PythonVersion.PY_312) +@register_handler("LOAD_ATTR", version=PythonVersion.PY_313) +@register_handler("LOAD_ATTR", version=PythonVersion.PY_314) +def handle_load_attr( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # LOAD_ATTR loads an attribute from the object on top of stack + obj = ensure_ast(state.stack[-1]) + attr_name = instr.argval + + # Create attribute access AST + attr_node = ast.Attribute(value=obj, attr=attr_name, ctx=ast.Load()) + if instr.argrepr.endswith(" + NULL|self"): + new_stack = state.stack[:-1] + [attr_node, Null()] + elif instr.argrepr.startswith("NULL|self + "): + new_stack = state.stack[:-1] + [Null(), attr_node] + else: + new_stack = state.stack[:-1] + [attr_node] + return replace(state, stack=new_stack) + + +@register_handler("BINARY_SUBSCR", version=PythonVersion.PY_312) +@register_handler("BINARY_SUBSCR", version=PythonVersion.PY_313) +def handle_binary_subscr( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # BINARY_SUBSCR implements obj[index] - pops index and obj from stack + index = ensure_ast(state.stack[-1]) # Index is on top + obj = ensure_ast(state.stack[-2]) # Object is below index + new_stack = state.stack[:-2] + + # Create subscript access AST + subscr_node = ast.Subscript(value=obj, slice=index, ctx=ast.Load()) + new_stack = new_stack + [subscr_node] + return replace(state, stack=new_stack) + + +@register_handler("BINARY_SLICE", version=PythonVersion.PY_312) +@register_handler("BINARY_SLICE", version=PythonVersion.PY_313) +@register_handler("BINARY_SLICE", version=PythonVersion.PY_314) +def handle_binary_slice( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # BINARY_SLICE implements obj[start:end] - pops start, end, and obj from stack + end = ensure_ast(state.stack[-1]) + start = ensure_ast(state.stack[-2]) + container = ensure_ast(state.stack[-3]) # Object is below start and end + sliced = ast.Subscript( + value=container, + slice=ast.Slice(lower=start, upper=end, step=None), + ctx=ast.Load(), + ) + new_stack = state.stack[:-3] + [sliced] + return replace(state, stack=new_stack) + + +# ============================================================================ +# OTHER CONTAINER BUILDING HANDLERS +# ============================================================================ + + +@register_handler("UNPACK_SEQUENCE", version=PythonVersion.PY_312) +@register_handler("UNPACK_SEQUENCE", version=PythonVersion.PY_313) +@register_handler("UNPACK_SEQUENCE", version=PythonVersion.PY_314) +def handle_unpack_sequence( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # UNPACK_SEQUENCE splits a comprehension loop target into `arg` sub-targets, + # as in ((k, v) for k, v in items). The names are not known yet, so the + # single target hole is refined into a tuple of fresh holes, which the + # following STORE_* instructions bind one at a time. + # + # CPython pushes the unpacked values right-to-left, so element 0 ends up on + # top of the stack and is consumed by the first STORE_*. + assert instr.arg is not None + unpack_count: int = instr.arg + + if not isinstance(state.stack[-1], TargetHole): + # Destructuring a known value rather than a loop target, as 3.14 emits + # when it unrolls a single-iteration loop over a literal. + elements = _literal_elements(state.stack[-1], unpack_count) + return replace(state, stack=state.stack[:-1] + list(reversed(elements))) + + holes = [TargetHole() for _ in range(unpack_count)] + new_stack = _bind_target_hole( + state.stack, state.stack[-1], ast.Tuple(elts=list(holes), ctx=ast.Store()) + ) + return replace(state, stack=new_stack[:-1] + list(reversed(holes))) + + +@register_handler("UNPACK_EX", version=PythonVersion.PY_312) +@register_handler("UNPACK_EX", version=PythonVersion.PY_313) +@register_handler("UNPACK_EX", version=PythonVersion.PY_314) +def handle_unpack_ex( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # UNPACK_EX handles a starred target, as in ((a, b) for a, *b in pairs). + # The low byte of the argument counts the targets before the starred one and + # the high byte counts those after it; the starred target itself collects + # whatever is left over. As with UNPACK_SEQUENCE the values are pushed + # right-to-left, so the first target ends up on top of the stack. + assert instr.arg is not None + before, after = instr.arg & 0xFF, instr.arg >> 8 + + if not isinstance(state.stack[-1], TargetHole): + # Destructuring a known value; the starred target collects the middle. + elements = _literal_elements(state.stack[-1]) + assert len(elements) >= before + after, "Too few values to unpack" + middle = elements[before : len(elements) - after] + unpacked: list[ast.expr] = [ + *elements[:before], + ast.List(elts=list(middle), ctx=ast.Load()), + *elements[len(elements) - after :], + ] + return replace(state, stack=state.stack[:-1] + list(reversed(unpacked))) + + holes = [TargetHole() for _ in range(before + 1 + after)] + elts: list[ast.expr] = list(holes) + elts[before] = ast.Starred(value=holes[before], ctx=ast.Store()) + + new_stack = _bind_target_hole( + state.stack, state.stack[-1], ast.Tuple(elts=elts, ctx=ast.Store()) + ) + return replace(state, stack=new_stack[:-1] + list(reversed(holes))) + + +@register_handler("BUILD_TUPLE", version=PythonVersion.PY_312) +@register_handler("BUILD_TUPLE", version=PythonVersion.PY_313) +@register_handler("BUILD_TUPLE", version=PythonVersion.PY_314) +def handle_build_tuple( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + assert instr.arg is not None + tuple_size: int = instr.arg + # Pop elements for the tuple + elements = ( + [ensure_ast(elem) for elem in state.stack[-tuple_size:]] + if tuple_size > 0 + else [] + ) + new_stack = state.stack[:-tuple_size] if tuple_size > 0 else state.stack + + # Create tuple AST + tuple_node = ast.Tuple(elts=elements, ctx=ast.Load()) + new_stack = new_stack + [tuple_node] + return replace(state, stack=new_stack) + + +@register_handler("BUILD_SLICE", version=PythonVersion.PY_312) +@register_handler("BUILD_SLICE", version=PythonVersion.PY_313) +@register_handler("BUILD_SLICE", version=PythonVersion.PY_314) +def handle_build_slice( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # BUILD_SLICE creates a slice object from the top of the stack + # The number of elements to pop is determined by the instruction argument + assert instr.arg is not None + slice_size: int = instr.arg + + if slice_size == 2: + # Slice with start and end: [start, end] + end = ensure_ast(state.stack[-1]) + start = ensure_ast(state.stack[-2]) + new_stack = state.stack[:-2] + slice_node = ast.Slice(lower=start, upper=end, step=None) + elif slice_size == 3: + # Slice with start, end, and step: [start, end, step] + step = ensure_ast(state.stack[-1]) + end = ensure_ast(state.stack[-2]) + start = ensure_ast(state.stack[-3]) + new_stack = state.stack[:-3] + slice_node = ast.Slice(lower=start, upper=end, step=step) + else: + raise ValueError(f"Unsupported slice size: {slice_size}") + + # Create slice AST + new_stack = new_stack + [slice_node] + return replace(state, stack=new_stack) + + +@register_handler("BUILD_CONST_KEY_MAP", version=PythonVersion.PY_312) +@register_handler("BUILD_CONST_KEY_MAP", version=PythonVersion.PY_313) +def handle_build_const_key_map( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # BUILD_CONST_KEY_MAP builds a dictionary with constant keys + # The keys are in a tuple on TOS, values are on the stack below + assert instr.arg is not None + assert isinstance(state.stack[-1], ast.Tuple), "Expected a tuple of keys" + map_size: int = instr.arg + # Pop the keys tuple and values + keys_tuple: ast.Tuple = state.stack[-1] + keys: list[ast.expr | None] = [ensure_ast(key) for key in keys_tuple.elts] + values = [ensure_ast(val) for val in state.stack[-map_size - 1 : -1]] + new_stack = state.stack[: -map_size - 1] + + # Create dictionary AST + dict_node = ast.Dict(keys=keys, values=values) + new_stack = new_stack + [dict_node] + return replace(state, stack=new_stack) + + +@register_handler("LIST_EXTEND", version=PythonVersion.PY_312) +@register_handler("LIST_EXTEND", version=PythonVersion.PY_313) +@register_handler("LIST_EXTEND", version=PythonVersion.PY_314) +def handle_list_extend( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # LIST_EXTEND appends the contents of the iterable at TOS to the list + # further down the stack. That list is either the empty ListComp that + # BUILD_LIST(0) optimistically created -- a list display, not a + # comprehension after all -- or a partly built argument list for a call + # with a starred argument. + update = state.stack[-1] + target = state.stack[-instr.argval - 1] + + # A literal iterable contributes its elements directly; anything else has to + # stay unpacked, as in `[*whatever]`. + elements: list[ast.expr] + if isinstance(update, ast.Tuple | ast.List): + elements = [ensure_ast(e) for e in update.elts] + else: + elements = [ast.Starred(value=ensure_ast(update), ctx=ast.Load())] + + if isinstance(target, ast.ListComp) and not target.generators: + merged = ast.List(elts=elements, ctx=ast.Load()) + else: + assert isinstance(target, ast.List), "LIST_EXTEND expects a list to extend" + merged = ast.List(elts=list(target.elts) + elements, ctx=ast.Load()) + + new_stack = state.stack[:-1] + new_stack[-instr.argval] = merged + return replace(state, stack=new_stack) + + +@register_handler("DICT_MERGE", version=PythonVersion.PY_312) +@register_handler("DICT_MERGE", version=PythonVersion.PY_313) +@register_handler("DICT_MERGE", version=PythonVersion.PY_314) +def handle_dict_merge( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # DICT_MERGE folds the mapping at TOS into the one below it, rejecting + # duplicate keys. It assembles the keyword arguments of a call using `**`. + update = state.stack[-1] + target = state.stack[-instr.argval - 1] + + # An `ast.Dict` entry with a key of None is `**value`, which is how a + # mapping that is not a literal has to be spliced in. + def entries(node: ast.expr) -> tuple[list[ast.expr | None], list[ast.expr]]: + if isinstance(node, ast.Dict): + return list(node.keys), list(node.values) + return [None], [ensure_ast(node)] + + if isinstance(target, ast.DictComp) and not target.generators: + # BUILD_MAP(0) guessed at a dict comprehension; it was a `**` argument. + keys, values = entries(update) + else: + assert isinstance(target, ast.Dict), "DICT_MERGE expects a dict to merge into" + target_keys, target_values = entries(target) + update_keys, update_values = entries(update) + keys, values = target_keys + update_keys, target_values + update_values + + new_stack = state.stack[:-1] + new_stack[-instr.argval] = ast.Dict(keys=keys, values=values) + return replace(state, stack=new_stack) + + +@register_handler("SET_UPDATE", version=PythonVersion.PY_312) +@register_handler("SET_UPDATE", version=PythonVersion.PY_313) +@register_handler("SET_UPDATE", version=PythonVersion.PY_314) +def handle_set_update( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # The set being extended is actually in state.result instead of the stack + # because it was initially recognized as a list comprehension in BUILD_SET, + # while the actual result expression is in the stack where the set "should be" + # and needs to be put back into the state result slot + assert isinstance(state.stack[-instr.argval - 1], ast.SetComp) + assert isinstance(state.stack[-1], ast.Tuple | ast.List | ast.Set) + + new_val = ast.Set(elts=[ensure_ast(e) for e in state.stack[-1].elts]) + new_stack = state.stack[:-2] + [new_val] + + return replace(state, stack=new_stack) + + +@register_handler("DICT_UPDATE", version=PythonVersion.PY_312) +@register_handler("DICT_UPDATE", version=PythonVersion.PY_313) +@register_handler("DICT_UPDATE", version=PythonVersion.PY_314) +def handle_dict_update( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # The dict being extended is actually in state.result instead of the stack + # because it was initially recognized as a list comprehension in BUILD_MAP, + # while the actual result expression is in the stack where the dict "should be" + # and needs to be put back into the state result slot + assert isinstance(state.stack[-instr.argval - 1], ast.DictComp) + assert isinstance(state.stack[-1], ast.Dict) + + new_val = ast.Dict( + keys=[ensure_ast(e) for e in state.stack[-1].keys], + values=[ensure_ast(e) for e in state.stack[-1].values], + ) + new_stack = state.stack[:-2] + [new_val] + + return replace(state, stack=new_stack) + + +@register_handler("BUILD_STRING", version=PythonVersion.PY_312) +@register_handler("BUILD_STRING", version=PythonVersion.PY_313) +@register_handler("BUILD_STRING", version=PythonVersion.PY_314) +def handle_build_string( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # BUILD_STRING concatenates strings from the stack + # For f-strings, it combines FormattedValue and Constant nodes + assert instr.arg is not None + string_size: int = instr.arg + + if string_size == 0: + # Empty string case + new_stack = state.stack + [ast.Constant(value="")] + return replace(state, stack=new_stack) + + # Pop elements for the string + elements = [ensure_ast(elem) for elem in state.stack[-string_size:]] + new_stack = state.stack[:-string_size] + + # Check if this is an f-string build (has FormattedValue nodes) + # or a regular string concatenation + if any(isinstance(elem, ast.JoinedStr) for elem in elements): + # This is an f-string - create JoinedStr + values = [] + for elem in elements: + if isinstance(elem, ast.JoinedStr): + values.extend(elem.values) + else: + values.append(elem) + return replace(state, stack=new_stack + [ast.JoinedStr(values=values)]) + elif all(isinstance(elem, ast.Constant) for elem in elements): + # This is regular string concatenation or format spec building + # If all elements are constants, we might be building a format spec + # Concatenate the constant strings + assert all( + isinstance(elem, ast.Constant) and isinstance(elem.value, str) + for elem in elements + ) + concat_str = "".join( + elem.value + for elem in elements + if isinstance(elem, ast.Constant) and isinstance(elem.value, str) + ) + return replace(state, stack=new_stack + [ast.Constant(value=concat_str)]) + else: + raise TypeError("Should not be here?") + + +@register_handler("FORMAT_VALUE", version=PythonVersion.PY_312) +def handle_format_value( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # FORMAT_VALUE formats a string with a value in Python 3.12 + # Flag bits: (flags & 0x03) = conversion, (flags & 0x04) = has format spec + assert instr.arg is not None, "FORMAT_VALUE requires flags argument" + assert len(state.stack) >= 1, "Not enough items on stack for FORMAT_VALUE" + + flags = instr.arg + + # Check if there's a format specification + has_format_spec = bool(flags & 0x04) + + if has_format_spec: + # Pop format spec and value + assert len(state.stack) >= 2, ( + "FORMAT_VALUE with format spec needs 2 stack items" + ) + format_spec = ensure_ast(state.stack[-1]) + value = ensure_ast(state.stack[-2]) + new_stack = state.stack[:-2] + + # Wrap format spec in JoinedStr if it's a constant + if isinstance(format_spec, ast.Constant): + format_spec_node = ast.JoinedStr(values=[format_spec]) + else: + assert isinstance(format_spec, ast.JoinedStr) + format_spec_node = format_spec + else: + # Just pop the value + value = ensure_ast(state.stack[-1]) + new_stack = state.stack[:-1] + format_spec_node = None + + # Determine conversion type from flags + conversion_flags = flags & 0x03 + conversion_map = { + 0: -1, # No conversion + 1: 115, # str (!s) + 2: 114, # repr (!r) + 3: 97, # ascii (!a) + } + conversion = conversion_map[conversion_flags] + + # Create formatted value AST + formatted_node = ast.FormattedValue( + value=value, conversion=conversion, format_spec=format_spec_node + ) + new_stack = new_stack + [ast.JoinedStr(values=[formatted_node])] + return replace(state, stack=new_stack) + + +@register_handler("FORMAT_SIMPLE", version=PythonVersion.PY_313) +@register_handler("FORMAT_SIMPLE", version=PythonVersion.PY_314) +def handle_format_simple( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # FORMAT_SIMPLE formats a string with a single value + # Pops the value and the format string from the stack + assert len(state.stack) >= 1, "Not enough items on stack for FORMAT_SIMPLE" + value = state.stack[-1] + + # Check if the value was converted + if isinstance(value, ConvertedValue): + conversion = value.ast_conversion + value = value.value + else: + conversion = -1 + value = ensure_ast(value) + + # Create formatted string AST + formatted_node = ast.FormattedValue( + value=value, conversion=conversion, format_spec=None + ) + new_stack = state.stack[:-1] + [ast.JoinedStr(values=[formatted_node])] + return replace(state, stack=new_stack) + + +@register_handler("FORMAT_WITH_SPEC", version=PythonVersion.PY_313) +@register_handler("FORMAT_WITH_SPEC", version=PythonVersion.PY_314) +def handle_format_with_spec( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # FORMAT_WITH_SPEC formats a value with a format specifier + # Stack order in Python 3.13: format_spec on top, value below + assert len(state.stack) >= 2, "Not enough items on stack for FORMAT_WITH_SPEC" + format_spec = ensure_ast(state.stack[-1]) # Format spec is on top + value = state.stack[-2] # Value is below + + # Check if the value was converted + if isinstance(value, ConvertedValue): + conversion = value.ast_conversion + value = value.value + else: + conversion = -1 + value = ensure_ast(value) + + # Create formatted string AST with specifier + # The format_spec should be wrapped in a JoinedStr if it's a simple constant + if isinstance(format_spec, ast.Constant): + format_spec_node = ast.JoinedStr(values=[format_spec]) + else: + # Already a JoinedStr from nested formatting + assert isinstance(format_spec, ast.JoinedStr) + format_spec_node = format_spec + + formatted_node = ast.FormattedValue( + value=value, conversion=conversion, format_spec=format_spec_node + ) + new_stack = state.stack[:-2] + [ast.JoinedStr(values=[formatted_node])] + return replace(state, stack=new_stack) + + +# ============================================================================ +# CONDITIONAL JUMP HANDLERS +# ============================================================================ + + +def _handle_pop_jump_if( + f_condition: Callable[[ast.expr], ast.expr], + state: ReconstructionState, + instr: dis.Instruction, +) -> ReconstructionState: + # Generic handler for POP_JUMP_IF_* instructions. Pops a value from the + # stack; `condition` is true exactly when the jump is taken. + condition: ast.expr = f_condition(ensure_ast(state.stack[-1])) + + # An inlined-builtin guard is an implementation detail of the interpreter, + # not part of the comprehension: drop it rather than record it as a filter. + if _specialization_guard_edge(state, instr) is not None: + return replace(state, stack=state.stack[:-1]) + + kind, _ = _classify_branch(state, instr) + edge = state.branches.get(instr.offset, BranchEdge.TAKE_JUMP) + + if kind is BranchKind.TERNARY: + return _handle_conditional_expression(state, instr, condition, edge) + + # A filter. The guard is the condition under which *this* path carries on + # toward the element, so it is negated when the path falls through. + guard = condition if edge is BranchEdge.TAKE_JUMP else _negate(condition) + return _attach_filter(state, guard) + + +def _attach_filter( + state: ReconstructionState, guard: ast.expr | None +) -> ReconstructionState: + """Conjoin ``guard`` to the filters of the innermost unfinished comprehension.""" + for pos, item in zip(reversed(range(len(state.stack))), reversed(state.stack)): + if not isinstance(item, CompExp): + continue + + elt: ast.expr = item.value if isinstance(item, ast.DictComp) else item.elt + new_result: CompExp = copy.deepcopy(item) + + if isinstance(elt, Placeholder): + resolved = guard + elif isinstance(elt, ast.IfExp) and any( + isinstance(x, Placeholder) for x in ast.walk(elt) + ): + # A conditional expression was being built up in the element slot, + # but it turned out to be part of this filter's condition. Move it + # back out, plugging the guard into the arm still awaiting a value. + if isinstance(new_result, ast.DictComp): + new_result.key, new_result.value = Placeholder(), Placeholder() + else: + new_result.elt = Placeholder() + + if guard is None: + resolved = None + else: + plugged = ReplacePlaceholder(guard).visit(copy.deepcopy(elt)) + assert isinstance(plugged, ast.expr) + resolved = plugged + else: + continue + + if resolved is not None: + ifs = new_result.generators[-1].ifs + combined = _conjoin(ifs + [resolved]) + assert combined is not None + new_result.generators[-1].ifs = [combined] + + new_stack = state.stack[:pos] + [new_result] + state.stack[pos + 1 : -1] + return replace(state, stack=new_stack) + + raise TypeError("No comprehension context found for filter condition") + + +def _handle_conditional_expression( + state: ReconstructionState, + instr: dis.Instruction, + condition: ast.expr, + edge: BranchEdge, +) -> ReconstructionState: + """Start an ``ast.IfExp``, marking the arm this path did not take.""" + for pos, item in zip(reversed(range(len(state.stack))), reversed(state.stack)): + if any(isinstance(x, Placeholder) for x in ast.walk(item)): + body: Skipped | Placeholder + orelse: Skipped | Placeholder + skipped = Skipped(id=f".SKIPPED_{instr.offset}") + if edge is BranchEdge.FALL_THROUGH: + body, orelse = skipped, Placeholder() + else: + body, orelse = Placeholder(), skipped + + new_ifexp = ast.IfExp(test=condition, body=body, orelse=orelse) + new_result = ReplacePlaceholder(new_ifexp).visit(copy.deepcopy(item)) + new_stack = state.stack[:pos] + [new_result] + state.stack[pos + 1 : -1] + return replace(state, stack=new_stack) + + raise TypeError("No placeholder found for conditional expression") + + +@register_handler("POP_JUMP_IF_TRUE", version=PythonVersion.PY_312) +@register_handler("POP_JUMP_IF_TRUE", version=PythonVersion.PY_313) +@register_handler("POP_JUMP_IF_TRUE", version=PythonVersion.PY_314) +def handle_pop_jump_if_true( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # POP_JUMP_IF_TRUE pops a value from the stack and jumps if it's true + # In Python 3.13, this is used for filter conditions where True means continue + return _handle_pop_jump_if(lambda c: c, state, instr) + + +@register_handler("POP_JUMP_IF_FALSE", version=PythonVersion.PY_312) +@register_handler("POP_JUMP_IF_FALSE", version=PythonVersion.PY_313) +@register_handler("POP_JUMP_IF_FALSE", version=PythonVersion.PY_314) +def handle_pop_jump_if_false( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # POP_JUMP_IF_FALSE pops a value from the stack and jumps if it's false + # In comprehensions, this is used for filter conditions + return _handle_pop_jump_if( + lambda c: ast.UnaryOp(op=ast.Not(), operand=c), state, instr + ) + + +@register_handler("POP_JUMP_IF_NONE", version=PythonVersion.PY_312) +@register_handler("POP_JUMP_IF_NONE", version=PythonVersion.PY_313) +@register_handler("POP_JUMP_IF_NONE", version=PythonVersion.PY_314) +def handle_pop_jump_if_none( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # POP_JUMP_IF_NONE pops a value and jumps if it's None + return _handle_pop_jump_if( + lambda c: ast.Compare( + left=c, ops=[ast.Is()], comparators=[ast.Constant(value=None)] + ), + state, + instr, + ) + + +@register_handler("POP_JUMP_IF_NOT_NONE", version=PythonVersion.PY_312) +@register_handler("POP_JUMP_IF_NOT_NONE", version=PythonVersion.PY_313) +@register_handler("POP_JUMP_IF_NOT_NONE", version=PythonVersion.PY_314) +def handle_pop_jump_if_not_none( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # POP_JUMP_IF_NOT_NONE pops a value and jumps if it's not None + return _handle_pop_jump_if( + lambda c: ast.Compare( + left=c, ops=[ast.IsNot()], comparators=[ast.Constant(value=None)] + ), + state, + instr, + ) + + +@register_handler("JUMP_FORWARD", version=PythonVersion.PY_312) +@register_handler("JUMP_FORWARD", version=PythonVersion.PY_313) +@register_handler("JUMP_FORWARD", version=PythonVersion.PY_314) +def handle_jump_forward( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # JUMP_FORWARD is used to jump forward in the code + # In generator expressions, this is often used to skip code in conditional logic + return state + + +@register_handler("JUMP_BACKWARD", version=PythonVersion.PY_312) +@register_handler("JUMP_BACKWARD", version=PythonVersion.PY_313) +@register_handler("JUMP_BACKWARD", version=PythonVersion.PY_314) +def handle_jump_backward( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # JUMP_BACKWARD is used to jump back to the beginning of a loop (replaces JUMP_ABSOLUTE in 3.13) + # In generator expressions, this typically indicates the end of the loop body + return state + + +@register_handler("JUMP_BACKWARD_NO_INTERRUPT", version=PythonVersion.PY_312) +@register_handler("JUMP_BACKWARD_NO_INTERRUPT", version=PythonVersion.PY_313) +@register_handler("JUMP_BACKWARD_NO_INTERRUPT", version=PythonVersion.PY_314) +def handle_jump_backward_no_interrupt( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + raise TypeError( + "JUMP_BACKWARD_NO_INTERRUPT instruction should not appear in generator comprehensions" + ) + + +@register_handler("JUMP_NO_INTERRUPT", version=PythonVersion.PY_312) +@register_handler("JUMP_NO_INTERRUPT", version=PythonVersion.PY_313) +@register_handler("JUMP_NO_INTERRUPT", version=PythonVersion.PY_314) +def handle_jump_no_interrupt( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + raise TypeError( + "JUMP_NO_INTERRUPT instruction should not appear in generator comprehensions" + ) + + +@register_handler("JUMP", version=PythonVersion.PY_312) +@register_handler("JUMP", version=PythonVersion.PY_313) +@register_handler("JUMP", version=PythonVersion.PY_314) +def handle_jump( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + raise TypeError("JUMP instruction should not appear in generator comprehensions") + + +@register_handler("EXTENDED_ARG", version=PythonVersion.PY_312) +@register_handler("EXTENDED_ARG", version=PythonVersion.PY_313) +@register_handler("EXTENDED_ARG", version=PythonVersion.PY_314) +def handle_extended_arg( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # EXTENDED_ARG prefixes an instruction whose argument does not fit in a + # byte. `dis` has already folded it into the following instruction's `arg`, + # so there is nothing left to do here. + return state + + +@register_handler("RESUME", version=PythonVersion.PY_312) +@register_handler("RESUME", version=PythonVersion.PY_313) +@register_handler("RESUME", version=PythonVersion.PY_314) +def handle_resume( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # RESUME is used for resuming execution after yield/await - mostly no-op for AST reconstruction + return state + + +@register_handler("SEND", version=PythonVersion.PY_312) +@register_handler("SEND", version=PythonVersion.PY_313) +@register_handler("SEND", version=PythonVersion.PY_314) +def handle_send( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + raise TypeError("SEND instruction should not appear in generator comprehensions") + + +# ============================================================================ +# UTILITY FUNCTIONS +# ============================================================================ + + +@functools.singledispatch +def ensure_ast(value) -> ast.expr: + """Ensure value is an AST node""" + raise TypeError(f"Cannot convert {type(value)} to AST node") + + +@ensure_ast.register +def _ensure_ast_ast(value: ast.expr) -> ast.expr: + """If already an AST node, return it as is""" + return value + + +@ensure_ast.register(int) +@ensure_ast.register(float) +@ensure_ast.register(str) +@ensure_ast.register(bytes) +@ensure_ast.register(bool) +@ensure_ast.register(complex) +@ensure_ast.register(type(None)) +def _ensure_ast_constant(value) -> ast.Constant: + return ast.Constant(value=value) + + +@ensure_ast.register +def _ensure_ast_tuple(value: tuple) -> ast.Tuple: + """Convert tuple to AST""" + return ast.Tuple(elts=[ensure_ast(v) for v in value], ctx=ast.Load()) + + +def _unconsumed(value: Iterator) -> typing.Any: + """Return the items an iterator has not yet yielded, as a concrete sequence. + + Built-in sequence iterators pickle as ``(iter, (underlying,), index)``, where + ``index`` is how far the iterator has advanced (absent or ``None`` when it + does not apply). ``reversed`` objects pickle with ``reversed`` as the + callable and count *down* from the end of the underlying sequence. + """ + reduced = value.__reduce__() + assert isinstance(reduced, tuple) and len(reduced) >= 2, ( + f"Cannot recover the contents of {type(value)}" + ) + if not reduced[1]: # an exhausted iterator pickles with no arguments + return () + + underlying = reduced[1][0] + index = reduced[2] if len(reduced) > 2 and reduced[2] is not None else 0 + return underlying[index::-1] if reduced[0] is reversed else underlying[index:] + + +@ensure_ast.register(type(iter((1,)))) +def _ensure_ast_tuple_iterator(value: Iterator) -> ast.Tuple: + return ensure_ast(tuple(_unconsumed(value))) # type: ignore + + +@ensure_ast.register +def _ensure_ast_list(value: list) -> ast.List: + return ast.List(elts=[ensure_ast(v) for v in value], ctx=ast.Load()) + + +@ensure_ast.register(type(iter([1]))) +@ensure_ast.register(type(iter({1: 2}.values()))) +@ensure_ast.register(type(iter({1: 2}.items()))) +@ensure_ast.register(type(iter(reversed([1])))) +@ensure_ast.register(reversed) +def _ensure_ast_list_iterator(value: Iterator) -> ast.List: + return ensure_ast(list(_unconsumed(value))) # type: ignore + + +@ensure_ast.register(type(iter("ab"))) # str_ascii_iterator +@ensure_ast.register(type(iter("\xe9b"))) # str_iterator +@ensure_ast.register(type(iter(b"ab"))) +@ensure_ast.register(type(iter(bytearray(b"ab")))) +def _ensure_ast_str_iterator(value: Iterator) -> ast.Constant: + remainder = _unconsumed(value) + # bytearray iteration yields ints, exactly as bytes iteration does + return ensure_ast( # type: ignore + bytes(remainder) if isinstance(remainder, bytearray) else remainder + ) + + +@ensure_ast.register(set) +@ensure_ast.register(frozenset) +def _ensure_ast_set(value: set | frozenset) -> ast.Set: + return ast.Set(elts=[ensure_ast(v) for v in value]) + + +@ensure_ast.register(type(iter({1}))) +def _ensure_ast_set_iterator(value: Iterator) -> ast.Set: + return ensure_ast(set(_unconsumed(value))) # type: ignore + + +@ensure_ast.register +def _ensure_ast_dict(value: dict) -> ast.Dict: + return ast.Dict( + keys=[ensure_ast(k) for k in value.keys()], + values=[ensure_ast(v) for v in value.values()], + ) + + +@ensure_ast.register(type(iter({1: 2}))) +def _ensure_ast_dict_iterator(value: Iterator) -> ast.expr: + return ensure_ast(_unconsumed(value)) + + +@ensure_ast.register(types.BuiltinFunctionType) +@ensure_ast.register(type) +def _ensure_ast_builtin(value: typing.Callable) -> ast.Name: + """A built-in callable is referred to by name, which resolves via builtins. + + Covers both built-in functions (``abs``) and built-in types used as + callables (``bool``), which appear as the predicate of a ``filter`` or the + function of a ``map``. + """ + name = getattr(value, "__name__", None) + assert name and getattr(builtins, name, None) is value, ( + f"Cannot reference non-builtin callable {value!r}" + ) + return ast.Name(id=name, ctx=ast.Load()) + + +@ensure_ast.register(zip) +@ensure_ast.register(enumerate) +@ensure_ast.register(map) +@ensure_ast.register(filter) +def _ensure_ast_iterator_adaptor(value: Iterator) -> ast.Call: + """Rebuild zip/enumerate/map/filter from the arguments they pickle with. + + These wrap other iterators rather than a concrete sequence, so unlike a list + or range iterator they cannot be materialised -- but ``__reduce__`` hands + back their constituent parts, each of which ``ensure_ast`` can handle in + turn. Any already-consumed prefix is reflected in the inner iterators. + + A ``zip`` also pickles its strictness as reduction state, which has to be + carried over as a keyword argument: a strict ``zip`` raises on ragged input + where a lax one stops at the shortest iterable. + """ + reduced = value.__reduce__() + if isinstance(reduced, str): + raise TypeError(f"Cannot convert {type(value)} to AST node") + func, args = reduced[:2] + keywords = [] + if isinstance(value, zip) and len(reduced) > 2 and reduced[2]: + keywords.append(ast.keyword(arg="strict", value=ast.Constant(value=True))) + return ast.Call( + func=ast.Name(id=func.__name__, ctx=ast.Load()), + args=[ensure_ast(arg) for arg in args], + keywords=keywords, + ) + + +@ensure_ast.register +def _ensure_ast_slice(value: slice) -> ast.Slice: + """A constant slice, as 3.14 emits for `s[1:3]` alongside BINARY_OP/NB_SUBSCR.""" + return ast.Slice( + lower=None if value.start is None else ensure_ast(value.start), + upper=None if value.stop is None else ensure_ast(value.stop), + step=None if value.step is None else ensure_ast(value.step), + ) + + +@ensure_ast.register +def _ensure_ast_range(value: range) -> ast.Call: + return ast.Call( + func=ast.Name(id="range", ctx=ast.Load()), + args=[ensure_ast(value.start), ensure_ast(value.stop), ensure_ast(value.step)], + keywords=[], + ) + + +@ensure_ast.register(type(iter(range(1)))) +def _ensure_ast_range_iterator(value: Iterator) -> ast.Call: + return ensure_ast(_unconsumed(value)) # type: ignore + + +def _cell_values( + code: types.CodeType, closure: tuple[types.CellType, ...] | None +) -> dict[str, typing.Any]: + """Read the cells a function closed over, by free-variable name. + + A cell that is still empty -- a recursive definition not yet bound, say -- + is left out. + """ + values: dict[str, typing.Any] = {} + for name, cell in zip(code.co_freevars, closure or ()): + try: + values[name] = cell.cell_contents + except ValueError: + continue + return values + + +def _freevar_bindings( + code: types.CodeType, captured: collections.abc.Mapping[str, typing.Any] +) -> dict[str, ast.expr]: + """AST nodes for the values ``code``'s free variables were captured from. + + A free variable with no captured value to be found is left out, and the + reconstructed tree keeps the bare name, which is as close as it can get. + + An iterator is refused rather than written in. `ensure_ast` spells one as + the elements it has left, which is what the outermost iterable wants -- it + is about to be consumed anyway -- but a captured iterator is a value the + body can do anything with, and a list of its remaining elements is not the + same object. + """ + bindings: dict[str, ast.expr] = {} + for name in code.co_freevars: + if name not in captured: + continue + value = captured[name] + try: + if isinstance(value, Iterator): + raise TypeError("an iterator has no AST spelling in value position") + bindings[name] = ensure_ast(value) + except (TypeError, AssertionError) as exc: + raise TypeError( + f"Cannot represent {value!r}, " + f"the value captured in free variable {name!r}: {exc}" + ) from exc + return bindings + + +def _reconstruct_code( + value: types.CodeType, freevars: dict[str, ast.expr] +) -> ast.Lambda | CompLambda: + """Reconstruct a lambda or comprehension body from its code object. + + ``freevars`` maps the names this code captures from enclosing scopes to the + values found in their cells; see `handle_load_deref`. + """ + assert inspect.iscode(value), "Input must be a code object" + + name: str = value.co_name.split(".")[-1] + + # Check preconditions + if name in {"", "", "", ""}: + assert name == "" or sys.version_info < (3, 13) + assert name != "" or value.co_flags & inspect.CO_GENERATOR + assert value.co_flags & inspect.CO_NEWLOCALS + assert value.co_argcount == 1 + assert value.co_kwonlyargcount == value.co_posonlyargcount == 0 + assert DummyIterName().id in value.co_varnames + elif name == "": + assert not value.co_flags & inspect.CO_GENERATOR + assert value.co_flags & inspect.CO_NEWLOCALS + assert DummyIterName().id not in value.co_varnames + else: + raise TypeError(f"Unsupported code object type: {value.co_name}") + + # Symbolic execution to reconstruct the AST + result: ast.expr = _symbolic_exec(value, freevars) + + # Check postconditions + assert not any(isinstance(x, ast.stmt) for x in ast.walk(result)), ( + "Final return value must not contain statement nodes" + ) + assert not any( + isinstance( + x, + Placeholder + | Skipped + | TargetHole + | CommonConstant + | Null + | CompLambda + | ConvertedValue, + ) + for x in ast.walk(result) + ), "Final return value must not contain temporary nodes" + assert not any(x.arg == ".0" for x in ast.walk(result) if isinstance(x, ast.arg)), ( + "Final return value must not contain .0 argument" + ) + assert not any( + isinstance(x, ast.Name) and x.id == ".0" + for x in ast.walk(result) + if not isinstance(x, DummyIterName) + ), "Final return value must not contain .0 names" + assert sum(1 for x in ast.walk(result) if isinstance(x, DummyIterName)) <= 1, ( + "Final return value must contain at most 1 dummy iterator names" + ) + assert all(x.generators for x in ast.walk(result) if isinstance(x, CompExp)), ( + "Return value must have generators if not a lambda" + ) + + if name == "" and isinstance(result, ast.expr): + # co_varnames lists parameters first: positional, keyword-only, then + # *args and **kwargs if present. Default values are not part of the code + # object -- they are pushed by the caller and attached by MAKE_FUNCTION + # (3.12) or SET_FUNCTION_ATTRIBUTE (3.13). + names = value.co_varnames + n_args, n_kwonly = value.co_argcount, value.co_kwonlyargcount + n_params = n_args + n_kwonly + + vararg = kwarg = None + if value.co_flags & inspect.CO_VARARGS: + vararg = ast.arg(arg=names[n_params]) + n_params += 1 + if value.co_flags & inspect.CO_VARKEYWORDS: + kwarg = ast.arg(arg=names[n_params]) + + args = ast.arguments( + posonlyargs=[ast.arg(arg=arg) for arg in names[: value.co_posonlyargcount]], + args=[ast.arg(arg=arg) for arg in names[value.co_posonlyargcount : n_args]], + vararg=vararg, + kwonlyargs=[ast.arg(arg=arg) for arg in names[n_args : n_args + n_kwonly]], + kw_defaults=[None] * n_kwonly, + kwarg=kwarg, + defaults=[], + ) + return ast.Lambda(args=args, body=result) + elif name == "" and isinstance(result, ast.GeneratorExp): + return CompLambda(body=result) + elif name == "" and isinstance(result, ast.DictComp): + return CompLambda(body=result) + elif name == "" and isinstance(result, ast.ListComp): + return CompLambda(body=result) + elif name == "" and isinstance(result, ast.SetComp): + return CompLambda(body=result) + else: + raise TypeError(f"Invalid result for type {name}: {result}") + + +@ensure_ast.register +def _ensure_ast_codeobj(value: types.CodeType) -> ast.Lambda | CompLambda: + """A bare code object has no cells attached, so free variables stay names.""" + return _reconstruct_code(value, {}) + + +@ensure_ast.register +def _ensure_ast_lambda(value: types.LambdaType) -> ast.Lambda: + assert inspect.isfunction(value) and value.__name__.endswith(""), ( + "Input must be a lambda function" + ) + code: types.CodeType = value.__code__ + result = _reconstruct_code( + code, _freevar_bindings(code, _cell_values(code, value.__closure__)) + ) + assert isinstance(result, ast.Lambda), "Lambda body must be an AST Lambda node" + assert not isinstance(result, CompLambda), "Lambda must not be a CompLambda" + + # Default values are not in the code object: they were evaluated where the + # lambda was written and attached to the function. A lambda built inside the + # comprehension gets them from the stack instead -- see + # `_apply_function_attribute` -- but one arriving as a live object carries + # them here, and dropping them would leave parameters with no way to be + # filled. They cover the *trailing* positional parameters. + if value.__defaults__: + result.args.defaults = [ensure_ast(d) for d in value.__defaults__] + if value.__kwdefaults__: + by_name = value.__kwdefaults__ + result.args.kw_defaults = [ + ensure_ast(by_name[arg.arg]) if arg.arg in by_name else None + for arg in result.args.kwonlyargs + ] + return result + + +@ensure_ast.register +def _ensure_ast_genexpr(genexpr: types.GeneratorType) -> ast.GeneratorExp: + assert inspect.isgenerator(genexpr), "Input must be a generator expression" + assert inspect.getgeneratorstate(genexpr) == inspect.GEN_CREATED, ( + "Generator must be in created state" + ) + assert genexpr.gi_frame is not None, "Generator must not be exhausted" + # A generator holds its cells in its frame rather than in a __closure__, and + # an unstarted frame already has them all copied in. + frame_locals = genexpr.gi_frame.f_locals + genexpr_ast = _reconstruct_code( + genexpr.gi_code, _freevar_bindings(genexpr.gi_code, frame_locals) + ) + assert isinstance(genexpr_ast, CompLambda) + geniter_ast = ensure_ast(frame_locals[".0"]) + result = genexpr_ast.inline(geniter_ast) + assert isinstance(result, ast.GeneratorExp) + assert inspect.getgeneratorstate(genexpr) == inspect.GEN_CREATED, ( + "Generator must stay in created state" + ) + return result + + +# ============================================================================ +# MAIN RECONSTRUCTION FUNCTION +# ============================================================================ + + +def disassemble( + genexpr: Generator[typing.Any, typing.Any, typing.Any], +) -> ast.Expression: + """ + Reconstruct an AST from a generator expression's bytecode. + + This function analyzes the bytecode of a generator object and reconstructs + an abstract syntax tree (AST) that represents the original comprehension + expression. The reconstruction process simulates the Python VM's execution + of the bytecode, building AST nodes instead of executing operations. + + The reconstruction handles complex comprehension features including: + - Multiple nested loops + - Filter conditions (if clauses) + - Complex expressions in the yield/result part + - Tuple unpacking in loop variables + - Various operators and function calls + + Args: + genexpr (Generator[object, None, None]): The generator object to analyze. + Must be a freshly created generator that has not been iterated yet + (in 'GEN_CREATED' state). + + Returns: + ast.Expression: An AST node representing the reconstructed comprehension. + + Raises: + ValueError: If the input is not a generator or if the generator + has already been started (not in 'GEN_CREATED' state). + TypeError: If some part of the comprehension has no AST spelling -- + an outermost iterable or a captured free variable whose value + cannot be written into the tree. + + Example: + >>> # Generator expression + >>> g = (x * 2 for x in range(10) if x % 2 == 0) + >>> ast_node = disassemble(g) + >>> isinstance(ast_node, ast.Expression) + True + + >>> # The reconstructed AST can be compiled and evaluated + >>> import ast + >>> code = compile(ast_node, '', 'eval') + >>> result = eval(code) + >>> list(result) + [0, 4, 8, 12, 16] + + Note: + The reconstruction is based on bytecode analysis and may not perfectly + preserve the original source code formatting or variable names in all + cases. However, the semantic behavior of the reconstructed AST should + match the original comprehension, subject to the limits set out under + "What is recovered, and what is not" in this module's docstring: + evaluating the result re-runs every expression in it, so a stateful + filter or element expression can answer differently the second time, + and the outermost iterable is a snapshot rather than an expression. + """ + if not inspect.isgenerator(genexpr): + raise ValueError( + f"Input must be a generator expression, got {type(genexpr).__name__}" + ) + if inspect.getgeneratorstate(genexpr) != inspect.GEN_CREATED: + raise ValueError( + "Input must be a generator expression that has not been started, " + f"got one in state {inspect.getgeneratorstate(genexpr)}" + ) + return ast.fix_missing_locations(ast.Expression(ensure_ast(genexpr))) diff --git a/tests/test_internals_disassembler.py b/tests/test_internals_disassembler.py new file mode 100644 index 000000000..22e484167 --- /dev/null +++ b/tests/test_internals_disassembler.py @@ -0,0 +1,2109 @@ +import ast +import collections.abc +import copy +import typing + +import pytest + +from effectful.internals.disassembly import ( + CompLambda, + DummyIterName, + disassemble, + ensure_ast, +) + + +def compile_and_eval( + node: ast.expr | ast.Expression, globals_dict: dict | None = None +) -> typing.Any: + """Compile an AST node and evaluate it.""" + if globals_dict is None: + globals_dict = {} + + # Wrap in an Expression node if needed + if not isinstance(node, ast.Expression): + node = ast.Expression(body=node) + + # Fix location info + ast.fix_missing_locations(node) + + # Compile and evaluate + code = compile(node, "", "eval") + return eval(code, globals_dict) + + +def materialize[T](genexpr: collections.abc.Generator[T, None, None]) -> list[T]: + """Materialize a nested generator expression to a nested list.""" + + def _materialize(genexpr): + if isinstance(genexpr, str | bytes): + return genexpr + elif isinstance(genexpr, collections.abc.Generator): + return [_materialize(item) for item in genexpr] + elif isinstance(genexpr, tuple): + # Kept as a tuple so that sets of tuples stay hashable + return tuple(_materialize(item) for item in genexpr) + elif isinstance(genexpr, collections.abc.Sequence): + return [_materialize(item) for item in genexpr] + elif isinstance(genexpr, collections.abc.Set): + return {_materialize(item) for item in genexpr} + elif isinstance(genexpr, collections.abc.Mapping): + return {_materialize(k): _materialize(v) for k, v in genexpr.items()} + else: + return genexpr + + return [_materialize(x) for x in genexpr] + + +def assert_ast_equivalent( + genexpr: collections.abc.Generator[typing.Any, None, None], + reconstructed_ast: ast.AST, + globals_dict: dict | None = None, +): + """Assert that a reconstructed AST produces the same results as the original generator.""" + # Check AST structure + assert isinstance(reconstructed_ast, ast.Expression) + assert hasattr(reconstructed_ast.body, "elt") # The expression part + assert hasattr(reconstructed_ast.body, "generators") # The comprehension part + assert len(reconstructed_ast.body.generators) > 0 + for comp in reconstructed_ast.body.generators: + assert hasattr(comp, "target") # Loop variable + assert hasattr(comp, "iter") # Iterator + assert hasattr(comp, "ifs") # Conditions + + # Save current globals to restore later + curr_globals = globals().copy() + globals().update(globals_dict or {}) + + # Materialize original generator to list for comparison + original_list = materialize(genexpr) + + # Clean up globals to avoid pollution + for key in globals_dict or {}: + if key not in curr_globals: + del globals()[key] + globals().update(curr_globals) + + # Compile and evaluate the reconstructed AST + reconstructed_gen = compile_and_eval(reconstructed_ast, globals_dict) + reconstructed_list = materialize(reconstructed_gen) + assert reconstructed_list == original_list, ( + f"AST produced {reconstructed_list}, expected {original_list}" + ) + + +# ============================================================================ +# BASIC GENERATOR EXPRESSION TESTS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # Simple generator expressions + (x for x in range(5)), + (y for y in range(10)), + (item for item in [1, 2, 3]), + # Edge cases for simple generators + (i for i in range(0)), # Empty range + (n for n in range(1)), # Single item range + (val for val in range(100)), # Large range + (x for x in range(-5, 5)), # Negative range + (step for step in range(0, 10, 2)), # Step range + (rev for rev in range(10, 0, -1)), # Reverse range + ], +) +def test_simple_generators(genexpr): + """Test reconstruction of simple generator expressions.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# ARITHMETIC AND EXPRESSION TESTS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # Basic arithmetic operations + (x * 2 for x in range(5)), + (x + 1 for x in range(5)), + (x - 1 for x in range(5)), + (x**2 for x in range(5)), + (x % 2 for x in range(10)), + (x / 2 for x in range(1, 6)), + (x // 2 for x in range(10)), + # Complex expressions + (x * 2 + 1 for x in range(5)), + ((x + 1) * (x - 1) for x in range(5)), + (x**2 + 2 * x + 1 for x in range(5)), + # Unary operations + (-x for x in range(5)), + (+x for x in range(-5, 5)), + (~x for x in range(5)), + # More complex arithmetic edge cases + (x**3 for x in range(1, 5)), # Higher powers + (x * x * x for x in range(5)), # Repeated multiplication + (x + x + x for x in range(5)), # Repeated addition + (x - x + 1 for x in range(5)), # Operations that might simplify + (x / x for x in range(1, 5)), # Division by self + (x % (x + 1) for x in range(1, 10)), # Modulo with expression + # Nested arithmetic expressions + ((x + 1) ** 2 for x in range(5)), + ((x * 2 + 3) * (x - 1) for x in range(5)), + (x * (x + 1) * (x + 2) for x in range(5)), + # Mixed operations with precedence + (x + 3 * 2 for x in range(3)), + (x * 2 + 9 / 3 for x in range(1, 4)), + ((x + 2) * (x - 2) for x in range(1, 4)), + # Edge cases with zero and one + (x * 0 for x in range(5)), + (x * 1 for x in range(5)), + (x + 0 for x in range(5)), + (x**1 for x in range(5)), + (0 + x for x in range(5)), + (1 * x for x in range(5)), + ], +) +def test_arithmetic_expressions(genexpr): + """Test reconstruction of generators with arithmetic expressions.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# FSTRING EXPRESSIONS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # Basic f-string cases + (f"{x}" for x in range(5)), # Single value, no format + (f"{x} is {x**2}" for x in range(5)), # Multiple values + (f"{x:02d}" for x in range(10)), # Format spec + (f"{x:.2f}" for x in [1.2345, 2.3456, 3.4567]), # Float format spec + # Conversion specifiers + (f"{x!r}" for x in ["hello", "world"]), # repr conversion + (f"{x!s}" for x in [1, 2, 3]), # str conversion + (f"{x!a}" for x in ["hello\n", "world\t"]), # ascii conversion + # Conversion with format spec + (f"{x!r:>10}" for x in ["hello", "world"]), # repr with alignment + (f"{x!s:^15}" for x in [1, 2, 3]), # str with center align + # Empty and literal f-strings + ("" for x in range(3)), # Empty f-string + ("constant" for x in range(3)), # No formatting + (f"x={x}" for x in range(5)), # Literal prefix + (f"result: {x * 2}" for x in range(5)), # Literal with expression + # Complex expressions in f-strings + (f"{x + 1}" for x in range(5)), # Arithmetic + (f"{x * x}" for x in range(5)), # Multiplication + (f"{x % 2}" for x in range(10)), # Modulo + (f"{-x}" for x in range(-2, 3)), # Unary minus + # Nested formatting + (f"{x:0{2}d}" for x in range(5)), # Format spec with expression + (f"{x:>{3 * 2}}" for x in range(5)), # Expression in format spec + # Multiple formatted values + (f"{x} + {y} = {x + y}" for x in range(3) for y in range(3)), # Multiple vars + (f"({x}, {y})" for x in range(2) for y in range(2)), # Tuple display + # F-strings with various data types + (f"{s}" for s in ["hello", "world"]), # Strings + (f"{b}" for b in [True, False]), # Booleans + (f"{n}" for n in [None, None]), # None values + (f"{lst}" for lst in [[1, 2], [3, 4]]), # Lists + # Complex format specifications + (f"{x:+05d}" for x in range(-2, 3)), # Sign, zero pad, width + (f"{x:.2%}" for x in [0.1, 0.25, 0.333]), # Percentage format + (f"{x:.2e}" for x in [100, 1000, 10000]), # Scientific notation + (f"{x:#x}" for x in [10, 15, 255]), # Hex with prefix + (f"{x:b}" for x in [2, 7, 15]), # Binary format + # Edge cases + ("{x}" for x in range(3)), # Escaped braces + (f"{{x}} = {x}" for x in range(3)), # Mixed escaped/formatted + (f"{{{x}}}" for x in range(3)), # Brace around formatted + ], +) +def test_fstring_expressions(genexpr): + """Test reconstruction of generators with f-string expressions.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# COMPARISON OPERATORS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # All comparison operators + (x for x in range(10) if x < 5), + (x for x in range(10) if x <= 5), + (x for x in range(10) if x > 5), + (x for x in range(10) if x >= 5), + (x for x in range(10) if x == 5), + (x for x in range(10) if x != 5), + # in/not in operators + (x for x in range(10) if x in [2, 4, 6, 8]), + (x for x in range(10) if x not in [2, 4, 6, 8]), + # is/is not operators (with None) + (x for x in [1, None, 3, None, 5] if x is not None), + (x for x in [1, None, 3, None, 5] if x is None), + # Boolean operations - these are complex cases that might need special handling + (x for x in range(10) if not x % 2), + (x for x in range(10) if not (x > 5)), + (x for x in range(10) if x > 2 and x < 8), + (x for x in range(10) if x < 3 or x > 7), + # More complex comparison edge cases + # Comparisons with expressions + (x for x in range(10) if x * 2 > 10), + (x for x in range(10) if x + 1 <= 5), + (x for x in range(10) if x**2 < 25), + (x for x in range(10) if (x + 1) * 2 != 6), + # Complex membership tests + (x for x in range(20) if x in range(5, 15)), + (x for x in range(10) if x not in range(3, 7)), + (x for x in range(10) if x % 2 in [0]), + (x for x in range(10) if x not in []), # Empty container + # Complex boolean combinations + (x for x in range(20) if not (x < 5 or x > 15)), + (x for x in range(20) if x > 5 and x < 15 and x % 2 == 0), + (x for x in range(20) if x < 5 or x > 15 or x == 10), + (x for x in range(20) if not (x > 5 and x < 15)), + # Mixed comparison and boolean operations + (x for x in range(20) if (x > 10 and x % 2 == 0) or (x < 5 and x % 3 == 0)), + (x for x in range(20) if not (x % 2 == 0 and x % 3 == 0)), + # Edge cases with identity comparisons + (x for x in [0, 1, 2, None, 4] if x is not None and x > 1), + (x for x in [True, False, 1, 0] if x is True), + (x for x in [True, False, 1, 0] if x is not False), + ], +) +def test_comparison_operators(genexpr): + """Test reconstruction of all comparison operators.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# CHAINED COMPARISON TESTS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # Chained comparisons + (x for x in range(20) if 5 < x < 15), + (x for x in range(20) if 0 <= x <= 10), + ], +) +def test_chained_comparison_operators(genexpr): + """Test reconstruction of chained (ternary) comparison operators.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# FILTERED GENERATOR TESTS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # Simple filters + (x for x in range(10) if x % 2 == 0), + (x for x in range(10) if x > 5), + (x for x in range(10) if x < 5), + (x for x in range(10) if x != 5), + # Complex filters + (x for x in range(20) if x % 2 == 0 if x % 3 == 0), + (x for x in range(100) if x > 10 if x < 90 if x % 5 == 0), + # Filters with expressions + (x * 2 for x in range(10) if x % 2 == 0), + (x**2 for x in range(10) if x > 3), + # Boolean operations in filters + (x for x in range(10) if not x % 2), + (x for x in range(10) if x > 2 and x < 8), + (x for x in range(10) if x < 3 or x > 7), + # More complex filter edge cases + (x for x in range(50) if x % 7 == 0), # Different modulo + (x for x in range(10) if x >= 0), # Always true condition + (x for x in range(10) if x < 0), # Always false condition + ( + x for x in range(20) if x % 2 == 0 and x % 3 == 0 + ), # Multiple conditions with and + ( + x for x in range(20) if x % 2 == 0 or x % 3 == 0 + ), # Multiple conditions with or + # Nested boolean operations + (x for x in range(20) if (x > 5 and x < 15) or x == 0), + (x for x in range(20) if not (x > 10 and x < 15)), + (x for x in range(50) if x > 10 and (x % 2 == 0 or x % 3 == 0)), + # Multiple consecutive filters + (x for x in range(100) if x > 20 if x < 80 if x % 10 == 0), + (x for x in range(50) if x % 2 == 0 if x % 3 != 0 if x > 10), + # Filters with complex expressions + (x + 1 for x in range(20) if (x * 2) % 3 == 0), + (x**2 for x in range(10) if x * (x + 1) > 10), + (x / 2 for x in range(1, 20) if x % (x // 2 + 1) == 0), + # Edge cases with truthiness + (x for x in range(10) if x), # Truthy filter + (x for x in range(-5, 5) if not x), # Falsy filter + (x for x in range(10) if bool(x % 2)), # Explicit bool conversion + ], +) +def test_filtered_generators(genexpr): + """Test reconstruction of generators with if conditions.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# NESTED LOOP TESTS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # Basic nested loops + ((x, y) for x in range(3) for y in range(3)), + (x + y for x in range(3) for y in range(3)), + (x * y for x in range(1, 4) for y in range(1, 4)), + # Nested with filters + ((x, y) for x in range(5) for y in range(5) if x < y), + (x + y for x in range(5) if x % 2 == 0 for y in range(5) if y % 2 == 1), + # Triple nested + (x + y + z for x in range(2) for y in range(3) for z in range(4)), + ((x, y, z) for x in range(2) for y in range(3) for z in range(4)), + # More complex nested loop edge cases + # Different sized ranges + ((x, y) for x in range(2) for y in range(5)), + ((x, y) for x in range(10) for y in range(2)), + # Asymmetric operations + (x - y for x in range(5) for y in range(3)), + (x / (y + 1) for x in range(1, 6) for y in range(3)), + (x**y for x in range(1, 4) for y in range(3)), + # Complex expressions with multiple variables + (x * y + x for x in range(3) for y in range(3)), + (x + y + x * y for x in range(1, 4) for y in range(1, 4)), + ((x + y) ** 2 for x in range(3) for y in range(3)), + # Filters on different loop levels + ((x, y) for x in range(10) if x % 2 == 0 for y in range(10) if y % 3 == 0), + (x * y for x in range(5) for y in range(5) if x != y), + (x + y for x in range(5) for y in range(5) if x + y < 5), + # Triple and quadruple nested with various patterns + (x + y + z for x in range(2) for y in range(2) for z in range(2)), + (x * y * z for x in range(1, 3) for y in range(1, 3) for z in range(1, 3)), + ( + (x, y, z, w) + for x in range(2) + for y in range(2) + for z in range(2) + for w in range(2) + ), + # Nested loops with complex filters + ((x, y) for x in range(5) if x > 1 for y in range(5) if x < y), + (x + y for x in range(3) if x > 0 for y in range(3)), + # Mixed range types + ((x, y) for x in range(-2, 2) for y in range(0, 4, 2)), + (x * y for x in range(5, 0, -1) for y in range(1, 6)), + # Dependent nested loops + ((x, y) for x in range(3) for y in range(x, 3)), + (x + y for x in range(3) for y in range(x + 1, 3)), + ], +) +def test_nested_loops(genexpr): + """Test reconstruction of generators with nested loops.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# =========================================================================== +# NESTED COMPREHENSIONS +# =========================================================================== + + +@pytest.mark.parametrize( + "genexpr", + [ + # nested generators + ((x for x in range(i + 1)) for i in range(5)), + ((x for j in range(i) for x in range(j)) for i in range(5)), + (((x for x in range(i + j)) for j in range(i)) for i in range(5)), + # nested generators with filters + ((x for x in range(i)) for i in range(5) if i > 0), + ((x for x in range(i) if x < i) for i in range(5) if i > 0), + (((x for x in range(i + j) if x < i + j) for j in range(i)) for i in range(5)), + # aggregation function call + (sum(x for x in range(i + 1)) for i in range(3)), + (max(x for x in range(i + 1)) for i in range(3)), + (dict((x, x + 1) for x in range(i + 1)) for i in range(3)), + (set(x for x in range(i + 1)) for i in range(3)), + # map + (list(map(abs, (x + 1 for x in range(i + 1)))) for i in range(3)), + (list(enumerate(x + 1 for x in range(i + 1))) for i in range(3)), + # nesting on both sides + ((y for y in range(x)) for x in (x_ + 1 for x_ in range(5))), + ((y for y in range(x)) for x in (x_ + 1 for x_ in range(5))), + ], +) +def test_nested_comprehensions(genexpr): + """Test reconstruction of nested comprehensions.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +def test_nested_comprehensions_multiline(): + """The same filter reconstructs the same way however the source is laid out. + + On Python 3.12 these two spellings disassemble to different jump layouts -- + only the one-line form emits POP_JUMP_IF_TRUE -- which used to make the + multiline form come out negated. + """ + one_line = (x for x in range(5) if x > 1) + assert_ast_equivalent(one_line, disassemble(one_line)) + + multiline = ( + x + for x in range(5) # comment to avoid reformatting + if x > 1 + ) + assert_ast_equivalent(multiline, disassemble(multiline)) + + assert ast.unparse(disassemble(x for x in range(5) if x > 1)) == ast.unparse( + disassemble( + x + for x in range(5) # comment to avoid reformatting + if x > 1 + ) + ) + + +# ============================================================================ +# DIFFERENT COMPREHENSION TYPES +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # Comprehensions as iterator constants + (x_ for x_ in [x for x in range(5)]), + (x_ for x_ in {x for x in range(5)}), + (x_ for x_ in {x: x**2 for x in range(5)}), + # Comprehensions as yield expressions + ([y * 2 for y in range(x + 1)] for x in range(3)), + ({y + 3 for y in range(x + 1)} for x in range(3)), + ({y: y**2 for y in range(x + 1)} for x in range(3)), + # nested non-generators + ([x for x in range(i)] for i in range(5)), + ([x for j in range(i) for x in range(j)] for i in range(5)), + ({x: x**2 for x in range(i)} for i in range(5)), + ([[x for x in range(i + j)] for j in range(i)] for i in range(5)), + # Nested comprehensions with filters inside + ([x for x in range(i)] for i in range(5) if i > 0), + ([x for x in range(i) if x < i] for i in range(5) if i > 0), + ([[x for x in range(i + j) if x < i + j] for j in range(i)] for i in range(5)), + ], +) +def test_different_comprehension_types(genexpr): + """Test reconstruction of different comprehension types.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# DICT DISPLAYS +# +# A dict display with dynamic keys is built by BUILD_MAP from key/value pairs +# that the compiler pushed in source order. Dicts compare equal whatever order +# they were built in, so these tests pin the order down directly: which of two +# equal keys wins, what `items()` yields, and when each subexpression runs. +# ============================================================================ + + +_EVAL_ORDER: list[str] = [] + + +def _note(tag: str, value: typing.Any) -> typing.Any: + """Record that this subexpression was evaluated, and pass its value along.""" + _EVAL_ORDER.append(tag) + return value + + +def test_dict_display_duplicate_keys(): + """The last of two equal keys wins, so the pairs must keep their order.""" + genexpr = ({x: "first", x: "second"} for x in range(1)) # noqa: F602 + reconstructed = disassemble(genexpr) + assert ast.unparse(reconstructed) == ( + "({x: 'first', x: 'second'} for x in range(0, 1, 1))" + ) + assert materialize(genexpr) == [{0: "second"}] + assert materialize(compile_and_eval(reconstructed)) == [{0: "second"}] + + +def test_dict_display_insertion_order(): + genexpr = ({x: "a", x + 1: "b", x + 2: "c"} for x in range(1)) + reconstructed = disassemble(genexpr) + expected = [[(0, "a"), (1, "b"), (2, "c")]] + assert [list(d.items()) for d in genexpr] == expected + assert [list(d.items()) for d in compile_and_eval(reconstructed)] == expected + + +def test_dict_display_evaluation_order(): + """Keys and values run left to right, key before its own value.""" + genexpr = ( + {_note("k1", "a"): _note("v1", 1), _note("k2", "b"): _note("v2", 2)} + for _ in range(1) + ) + reconstructed = disassemble(genexpr) + + _EVAL_ORDER.clear() + assert materialize(genexpr) == [{"a": 1, "b": 2}] + assert _EVAL_ORDER == ["k1", "v1", "k2", "v2"] + + _EVAL_ORDER.clear() + assert materialize(compile_and_eval(reconstructed, {"_note": _note})) == [ + {"a": 1, "b": 2} + ] + assert _EVAL_ORDER == ["k1", "v1", "k2", "v2"] + + +# ============================================================================ +# CONDITIONAL EXPRESSIONS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # simple conditional expressions without nesting + ((lambda x: x if x % 2 == 0 else -x)(xi) for xi in range(5)), + ((lambda x: (x + 1) if x < 5 else (x - 1))(xi) for xi in range(10)), + ((lambda x: (x * 2) if x > 0 else (x / 2))(xi) for xi in range(-5, 5)), + ((lambda x: (x**2) if x != 0 else 1)(xi) for xi in range(-3, 4)), + # simple conditional expressions with negation + ((lambda x: (x + 10) if not (x < 5) else (x - 10))(xi) for xi in range(20)), + ((lambda x: (x * 3) if not (x % 2 == 0) else (x // 3))(xi) for xi in range(10)), + ((lambda x: (x**3) if not (x < 0) else (x**0.5))(xi) for xi in range(-5, 15)), + # conditional expressions with lazy test + ( + (lambda x: (x + 10) if (x > 5 and x < 15) else (x - 10))(xi) + for xi in range(20) + ), + ( + (lambda x: (x * 3) if (x % 2 == 0 or x % 3 == 0) else (x // 3))(xi) + for xi in range(10) + ), + ( + (lambda x: (x**3) if not (x < 0 or x > 10) else (x**0.5))(xi) + for xi in range(-5, 15) + ), + ], +) +def test_conditional_expressions_simple_no_comprehension(genexpr): + """Test reconstruction of simple conditional expressions isolated from comprehension bodies.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +@pytest.mark.parametrize( + "genexpr", + [ + # nested conditional expressions + ( + (lambda x: (x + 1) if x < 5 else ((x - 1) if x < 10 else (x * 2)))(xi) + for xi in range(15) + ), + ( + ( + lambda x: ( + (x * 2) if x % 2 == 0 else ((x // 2) if x % 3 == 0 else (x + 2)) + ) + )(xi) + for xi in range(10) + ), + ( + (lambda x: (x**2) if x > 0 else ((-x) ** 2 if x < -5 else 1))(xi) + for xi in range(-10, 5) + ), + ], +) +def test_conditional_expressions_nested_no_comprehension(genexpr): + """Test reconstruction of nested conditional expressions isolated from comprehension bodies.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +@pytest.mark.parametrize( + "genexpr", + [ + # Basic conditional expressions in comprehension bodies + ((x if x % 2 == 0 else -x) for x in range(5)), + ((x * 2 if x > 0 else x / 2) for x in range(-3, 4)), + ((x**2 if x != 0 else 1) for x in range(-2, 3)), + # Conditional expressions with filters + ((x if x % 2 == 0 else -x) for x in range(10) if x > 2), + ((x * 3 if x > 5 else x + 1) for x in range(20) if x % 3 == 0), + # Nested loops with conditional expressions + ((x + y if x > y else x - y) for x in range(3) for y in range(3)), + ( + (x * y if x != 0 and y != 0 else 0) + for x in range(-2, 3) + for y in range(-2, 3) + ), + # Multiple conditional expressions + ( + (x if x > 0 else 0) + (y if y > 0 else 0) + for x in range(-2, 3) + for y in range(-2, 3) + ), + # Conditional expressions in different parts + ([x if x > 0 else -x for x in range(i)] for i in range(1, 4)), + ((x if x % 2 == 0 else -x) for x in (y if y > 2 else y + 10 for y in range(5))), + # Complex nested conditional expressions + ((x if x > 0 else (x + 5 if x > -3 else x * 2)) for x in range(-5, 5)), + ((x * 2 if x > 0 else (x / 2 if x < 0 else 1)) for x in range(-3, 4)), + # Conditional expressions with function calls + ((abs(x) if x < 0 else x) for x in range(-3, 4)), + ((max(x, 0) if x is not None else 0) for x in [None, -1, 0, 1, 2]), + # Mixed with other complex expressions + ((x + 1 if x % 2 == 0 else x - 1) * 2 for x in range(5)), + ((x, y, x + y if x > y else x - y) for x in range(3) for y in range(3)), + ], +) +def test_conditional_expressions_simple_comprehensions(genexpr): + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +@pytest.mark.parametrize( + "genexpr", + [ + # `and` chains compile to a run of jumps that each fall through to the + # loop back-edge, which the disassembler folds into a single filter. + (x for x in range(10) if x > 2 and x < 8), + (x for x in range(20) if x > 5 and x % 2 == 0 and x < 15), + (x for x in range(-10, 10) if abs(x) > 3 and x % 2 == 0), + (x for x in ["hello", "world", "test"] if len(x) > 3 and x.startswith("h")), + (x for x in range(20) if x % 2 == 0 and x % 3 == 0 and x > 0 and x < 18), + ((x, y) for x in range(5) for y in range(5) if x < y and x + y > 2), + # `or` in filter position: reconstructed incorrectly, see the marker. + (x for x in range(10) if x < 3 or x > 7), + (x for x in range(20) if x < 5 or x > 15 or x == 10), + (x for x in range(20) if (x > 10 and x % 2 == 0) or (x < 5 and x % 3 == 0)), + (x for x in range(20) if x > 5 and (x < 10 or x > 15)), + (x for x in range(100) if (x > 10 and x < 50) and (x % 3 == 0 or x % 5 == 0)), + # `not (a and b)` is compiled exactly like `not a or not b`. + (x for x in range(100) if not (x > 30 and x < 70)), + # Chained comparisons in filter position. + (x for x in range(20) if 5 < x < 15), + (x for x in range(20) if 0 <= x <= 10), + (x for x in range(50) if 10 < x < 20 < x * 2), + (x for x in range(10) if 0 <= x <= 5 <= x + 5), + (x for x in range(50) if 5 < x < 15 and x % 2 == 0), + (x for x in range(50) if x > 20 or 5 < x < 15), + ], +) +def test_lazy_boolean_and_chained_comparisons_in_filters(genexpr): + """Lazy boolean operators and chained comparisons in *filter* position. + + This is the hard case: a filter's condition is recognised structurally, by + the jump falling through to the loop back-edge, so any condition CPython + compiles with an intermediate join point is misread. + """ + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +@pytest.mark.parametrize( + "genexpr", + [ + # The same operators in *ternary* position all work: both arms produce a + # value, so the fork/merge machinery in _symbolic_exec applies directly. + ((x if x > 5 and x < 15 else 0) for x in range(20)), + ((x if x < 3 or x > 17 else -x) for x in range(20)), + ((x if 5 < x < 15 else 0) for x in range(20)), + ((x * 2 if 0 <= x <= 10 else x / 2) for x in range(-5, 15)), + ((x if x > 2 and x < 8 else -x) for x in range(10)), + ((x if x < 2 or x > 8 else -x) for x in range(10)), + ((x if not (x > 2 and x < 8) else -x) for x in range(10)), + ((x if 0 <= x <= 5 <= x + 5 else -x) for x in range(10)), + ((x if x > 1 and x < 9 or x == 0 else -x) for x in range(10)), + # ... including nested inside another ternary + ((x if x > 5 or x < 2 else (0 if x == 3 else 1)) for x in range(10)), + ((x if x > 5 else (0 if 2 < x < 4 else 1)) for x in range(10)), + ], +) +def test_lazy_boolean_and_chained_comparisons_in_ternaries(genexpr): + """Lazy boolean operators and chained comparisons in conditional expressions.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +def test_short_circuit_filter_is_a_disjunction_of_paths(): + """A short-circuiting filter is rebuilt as the disjunction over its paths. + + `x < 3 or x > 7` used to be misread as a conditional expression, yielding an + AST that compiled and ran but computed `[0, 1, 2, False, False, ...]` instead + of `[0, 1, 2, 8, 9]`. Each path through the condition now contributes one + disjunct, so the reconstruction is equivalent rather than merely plausible. + """ + genexpr = (x for x in range(10) if x < 3 or x > 7) + reconstructed = disassemble(genexpr) + + assert isinstance(reconstructed.body, ast.GeneratorExp) + filters = reconstructed.body.generators[0].ifs + assert len(filters) == 1 + assert isinstance(filters[0], ast.BoolOp) and isinstance(filters[0].op, ast.Or) + assert materialize(compile_and_eval(reconstructed)) == [0, 1, 2, 8, 9] + + +@pytest.mark.parametrize( + "genexpr", + [ + # Simple conditional as function argument + (max(x if x > 0 else 0, 1) for x in range(-2, 3)), + (abs(x if x < 0 else -x) for x in range(-3, 3)), + (len(str(x) if x > 10 else "small") for x in range(15)), + # Multiple conditional arguments + ( + max(x if x > 0 else 0, y if y > 0 else 0) + for x in range(-1, 2) + for y in range(-1, 2) + ), + ( + pow(x if x != 0 else 1, y if y > 0 else 1) + for x in range(3) + for y in range(3) + ), + # Nested function calls with conditionals + (max(abs(x if x < 0 else -x), 1) for x in range(-3, 4)), + (int(str(x if x > 5 else x + 10)) for x in range(10)), + # Conditionals in keyword arguments (using dict constructor as example) + (dict(a=x if x > 0 else 0, b=x * 2 if x < 5 else x) for x in range(8)), + # Method calls with conditional arguments + ([1, 2, 3].index(x if x in [1, 2, 3] else 1) for x in range(5)), + ("hello".replace("l", x if isinstance(x, str) else "X") for x in ["a", 1, "b"]), + # Complex nested case: conditional in function argument, function call in conditional + (abs(x if len(str(x)) > 1 else x * 10) for x in range(15)), + # Mixed: conditional in function call within comprehension filter + (x for x in range(20) if max(x if x > 10 else 0, 5) > 8), + ], +) +def test_conditional_expressions_function_arguments(genexpr): + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# GENERATOR EXPRESSION WITH GLOBALS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr,globals_dict", + [ + # Using constants + ((x + a for x in range(5)), {"a": 10}), # type: ignore # noqa: F821 + ((data[i] for i in range(2)), {"data": [3, 4]}), # type: ignore # noqa: F821 + # Using global functions + ((abs(x) for x in range(-5, 5)), {"abs": abs}), + ((len(s) for s in ["a", "ab", "abc"]), {"len": len}), + ((max(x, 5) for x in range(10)), {"max": max}), + ((min(x, 5) for x in range(10)), {"min": min}), + ((round(x / 3, 2) for x in range(10)), {"round": round}), + ], +) +def test_variable_lookup(genexpr, globals_dict): + """Test reconstruction of expressions with globals.""" + ast_node = disassemble(genexpr) + + # Need to provide the same globals for evaluation + assert_ast_equivalent(genexpr, ast_node, globals_dict) + + +# ============================================================================ +# EDGE CASES AND COMPLEX SCENARIOS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr,globals_dict", + [ + # Using lambdas and functions + (((lambda y: y * 2)(x) for x in range(5)), {}), + (((lambda y: y + 1)(x) for x in range(5)), {}), + (((lambda y: y**2)(x) for x in range(5)), {}), + (((lambda a, b: a + b)(x, x) for x in range(5)), {}), + (((lambda: (x for x in range(i)))() for i in range(3)), {}), + ((f(x) for x in range(5)), {"f": lambda y: y * 3}), # type: ignore # noqa: F821 + # Attribute access + ((x.real for x in [1 + 2j, 3 + 4j, 5 + 6j]), {}), + ((x.imag for x in [1 + 2j, 3 + 4j, 5 + 6j]), {}), + ((x.conjugate() for x in [1 + 2j, 3 + 4j, 5 + 6j]), {}), + # slicing and indexing + ((s[:2] for s in ["hello", "world"]), {}), + ((s[1:3] for s in ["hello", "world"]), {}), + ((s[-1] for s in ["hello", "world"]), {}), + ((s[0:3] for s in ["hello", "world"]), {}), + ((s[::-1] for s in ["hello", "world"]), {}), + ((s[1:2:] for s in ["hello", "world"]), {}), + # Method calls + ((s.upper() for s in ["hello", "world"]), {}), + ((s.lower() for s in ["HELLO", "WORLD"]), {}), + ((s.strip() for s in [" hello ", " world "]), {}), + ((x.bit_length() for x in range(1, 10)), {}), + ((str(x).zfill(3) for x in range(10)), {"str": str}), + # Subscript operations + (((10, 20, 30)[i] for i in range(3)), {}), + (([10, 20, 30][i] for i in range(3)), {}), + (({"a": 1, "b": 2, "c": 3}[k] for k in ["a", "b", "c"]), {}), + (("hello"[i] for i in range(5)), {}), + ((data[i][j] for i in range(2) for j in range(2)), {"data": [[1, 2], [3, 4]]}), # type: ignore # noqa: F821 + # # More complex attribute chains + # ((obj.value.bit_length() for obj in [type('', (), {'value': x})() for x in range(1, 5)]), {}), + # Multiple function calls + ((abs(max(x, -x)) for x in range(-3, 4)), {"abs": abs, "max": max}), + ((len(str(x)) for x in range(100, 110)), {"len": len, "str": str}), + # Mixed operations + ( + (abs(x) + len(str(x)) for x in range(-10, 10)), + {"abs": abs, "len": len, "str": str}, + ), + ((s.upper().lower() for s in ["Hello", "World"]), {}), + # Edge cases with complex data structures + (((1, 2, 3)[x % 3] for x in range(10)), {}), + (([1, 2, 3][x % 3] for x in range(10)), {}), + (({1, 2, 3} for x in range(10)), {}), + # (({"even": x, "odd": x + 1}["even" if x % 2 == 0 else "odd"] for x in range(5)), {}), + # Function calls with multiple arguments + ((pow(x, 2, 10) for x in range(5)), {"pow": pow}), + ((divmod(x, 3) for x in range(10)), {"divmod": divmod}), + ], +) +def test_complex_scenarios(genexpr, globals_dict): + """Test reconstruction of complex generator expressions.""" + ast_node = disassemble(genexpr) + + # Need to provide the same globals for evaluation + assert_ast_equivalent(genexpr, ast_node, globals_dict) + + +# ============================================================================ +# UNPACKING LOOP TARGETS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # Simple tuple targets + ((a, b) for a, b in [(1, 2), (3, 4)]), + (a + b for a, b in [(1, 2), (3, 4)]), + (a * b for a, b in [(2, 3), (4, 5)]), + ((b, a) for a, b in [(1, 2), (3, 4)]), + ((a, b, c) for a, b, c in [(1, 2, 3), (4, 5, 6)]), + ((a, b, c, d) for a, b, c, d in [(1, 2, 3, 4)]), + # Nested tuple targets + ((a, b, c) for a, (b, c) in [(1, (2, 3)), (4, (5, 6))]), + ((a, b, c) for (a, b), c in [((1, 2), 3)]), + ((a, b, c, d) for (a, b), (c, d) in [((1, 2), (3, 4))]), + ((a, b, c) for a, (b, (c,)) in [(1, (2, (3,)))]), + # Unpacking over dict views + ((k, v) for k, v in {1: "a", 2: "b"}.items()), + (v for k, v in {1: "a", 2: "b"}.items()), + # Unpacking combined with filters + ((a, b) for a, b in [(1, 2), (3, 1)] if a < b), + (a + b for a, b in [(1, 2), (3, 4)] if a % 2 == 0), + ((a, b) for a, b in [(1, 2), (3, 4)] if a > 0 if b > 3), + # Unpacking in nested loops, in either position + ((x, a, b) for x in range(2) for a, b in [(1, 2), (3, 4)]), + ((a, b, y) for a, b in [(1, 2)] for y in range(2)), + ((a, b, c, d) for a, b in [(1, 2)] for c, d in [(3, 4), (5, 6)]), + ((a, b, x) for a, b in [(1, 2), (3, 4)] for x in range(a)), + # Unpacking inside other comprehension types + ([a for a, b in [(1, 2), (3, 4)]] for _ in range(2)), + ({a for a, b in [(1, 2), (3, 4)]} for _ in range(2)), + ({a: b for a, b in [(1, 2), (3, 4)]} for _ in range(2)), + ((a for a, b in [(1, 2), (3, 4)]) for _ in range(2)), + # Unpacking with a conditional expression in the body + ((a if a > b else b) for a, b in [(1, 2), (4, 3)]), + ], +) +def test_unpacking_targets(genexpr): + """Test reconstruction of comprehensions that unpack their loop target.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +@pytest.mark.parametrize( + "genexpr", + [ + ((x, a, b) for x in range(2) for a, b in [(1, 2)]), + ((a, b, c, d) for a, b in [(1, 2)] for c, d in [(3, 4)]), + ((x, a, b) for x in range(2) for a, *b in [(1, 2, 3)]), + ((x, a, b, c) for x in range(2) for a, (b, c) in [(1, (2, 3))]), + ((x, a) for x in range(2) for (a,) in [(1,)]), + ((x, a, b) for x in range(2) for a, *b in [(1,)]), # type: ignore[var-annotated] + ((x, y, a) for x in range(2) for y in range(2) for a, b in [(1, 2)]), + ], +) +def test_unpacking_over_single_element_literal(genexpr): + """A one-element inner loop over a literal. + + Python 3.14 unrolls this: it assigns the targets outright and emits no + FOR_ITER, so the loop is not there to be recovered. The names it bound are + substituted at their uses instead, which reproduces the same elements. + """ + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +@pytest.mark.parametrize( + "genexpr", + [ + # Starred last: UNPACK_EX with only a "before" count + (a for a, *b in [(1, 2, 3), (4, 5, 6)]), + (b for a, *b in [(1, 2, 3), (4, 5, 6)]), + ((a, b) for a, *b in [(1, 2, 3)]), + ((a, b, c) for a, b, *c in [(1, 2, 3, 4)]), + # Starred first: the "after" count lives in the high byte of the + # argument, so the instruction is prefixed with EXTENDED_ARG + (a for *a, b in [(1, 2, 3), (4, 5, 6)]), + (b for *a, b in [(1, 2, 3), (4, 5, 6)]), + ((a, b) for *a, b in [(1, 2, 3)]), + # Starred in the middle + ((a, b, c) for a, *b, c in [(1, 2, 3, 4)]), + ((a, b, c) for a, *b, c in [(1, 2, 3, 4, 5)]), + # Starred target that collects nothing + ((a, b) for a, *b in [(1,)]), # type: ignore[var-annotated] + # Combined with filters, nesting and other comprehension types + ((a, b) for a, *b in [(1, 2), (3, 4)] if a > 1), + ((x, a, b) for x in range(2) for a, *b in [(1, 2, 3), (4, 5, 6)]), + ([a for a, *b in [(1, 2, 3)]] for _ in range(2)), + ], +) +def test_unpacking_starred_targets(genexpr): + """Test reconstruction of starred loop targets (UNPACK_EX).""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# OUTERMOST ITERABLE TYPES +# +# The outermost iterable is not part of the comprehension's bytecode: it is a +# live object reachable through `gi_frame.f_locals[".0"]`, so `ensure_ast` has +# to rebuild an expression for it from the object alone. +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # Strings and bytes + (c for c in "hello"), + (c.upper() for c in "hello" if c != "l"), + (c for c in "h\xe9llo"), # non-ASCII takes a different iterator type + (b for b in b"abc"), + (b for b in bytearray(b"abc")), + # Sequences + (x for x in [1, 2, 3]), + (x for x in (1, 2, 3)), + (x for x in range(3)), + # Sets and frozensets + (x for x in {1, 2, 3}), + (x for x in frozenset({1, 2, 3})), + # Dict views + (k for k in {1: "a", 2: "b"}), + (k for k in {1: "a", 2: "b"}.keys()), + (v for v in {1: "a", 2: "b"}.values()), + (kv for kv in {1: "a", 2: "b"}.items()), + # reversed() over each of the underlying sequence types + (x for x in reversed([1, 2, 3])), + (x for x in reversed((1, 2, 3))), + (c for c in reversed("abc")), + (x for x in reversed(range(3))), + # Comprehensions as the outermost iterable + (x for x in (y for y in range(3))), + (x for x in [y for y in range(3)]), + (x for x in {y for y in range(3)}), + # Nested/structured contents + (t for t in [(1, 2), (3, 4)]), + (d for d in [{"a": 1}, {"b": 2}]), + (x for x in [[1, 2], [3, 4]]), + ], +) +def test_outermost_iterable_types(genexpr): + """Test reconstruction of the outermost iterable from the live object.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +@pytest.mark.parametrize( + "genexpr", + [ + # zip/enumerate/map/filter wrap other iterators rather than a concrete + # sequence, but pickle with their constituent parts. + (x for x in zip([1, 2], [3, 4])), + ((a, b) for a, b in zip([1, 2], [3, 4])), + (x for x in zip("ab", range(2), [7, 8])), + (x for x in enumerate("ab")), + ((i, c) for i, c in enumerate("abc")), + (x for x in enumerate("ab", 5)), + (x for x in map(abs, [-1, 2])), + (x for x in map(max, [1, 2], [3, 0])), + (x for x in filter(None, [0, 1, 2])), + (x for x in filter(bool, [0, 1, 2])), + # A strict zip over equal-length iterables behaves like a lax one, but + # its strictness has to survive anyway -- see the ragged case below. + ((a, b) for a, b in zip([1, 2], [3, 4], strict=True)), + (x for x in zip("ab", range(2), [7, 8], strict=True)), + (x for x in zip(range(2), map(abs, [-1, -2]), strict=True)), + # Nested adaptors, and adaptors over non-sequence iterables + (x for x in zip(range(2), map(abs, [-1, -2]))), + (x for x in enumerate(filter(None, [0, 1]))), + (x for x in map(abs, range(-2, 2))), + (x for x in zip("ab", (y for y in range(2)))), + # With a filter and a non-trivial element expression + (a + b for a, b in zip([1, 2], [3, 4]) if a > 1), + ], +) +def test_outermost_iterable_adaptors(genexpr): + """zip/enumerate/map/filter are rebuilt from the parts they pickle with.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +@pytest.mark.parametrize( + "genexpr", + [ + # A conditional expression as an inner loop's iterable. Its value is + # consumed by FOR_ITER rather than by the yield, so the half-built + # IfExp has to be pulled back out of the element slot. + (y for x in range(4) for y in (range(x) if x % 2 == 0 else range(x, x + 2))), + (y for x in range(4) for y in ([x] if x else [0])), + (y for x in range(4) for y in (range(x) if x > 1 else range(1))), + # An *empty* list literal as an arm is misread: BUILD_LIST(0) is also + # how an inlined list comprehension starts, and nothing later + # disambiguates the two here. + pytest.param( + (y for x in range(4) for y in ([x] if x else [])), + marks=pytest.mark.xfail( + strict=True, + reason="an empty list literal is indistinguishable from the start of a list comprehension", + ), + ), + ((x, y) for x in range(3) for y in ([0] if x % 2 else [1, 2])), + # ... with a filter on the inner loop, and nested two deep + (y for x in range(4) for y in (range(x) if x % 2 == 0 else [9]) if y > 0), + pytest.param( + ( + z + for x in range(3) + for y in (range(x) if x else [0]) + for z in ([y] if y else [7]) + ), + marks=pytest.mark.xfail( + strict=True, + reason="two conditional iterables in one comprehension leave paths that do not pairwise merge", + ), + ), + # ... and one whose arms are comprehensions of different kinds + (y for x in range(3) for y in ([i for i in range(x)] if x else {8})), + ], +) +def test_conditional_expression_as_iterable(genexpr): + """Test a conditional expression in the iterable position of a for clause.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +@pytest.mark.parametrize( + "genexpr", + [ + # An always-false filter lets the compiler drop the body entirely, so no + # element expression survives in the bytecode. The loop still runs. + (x for x in range(6) if False), + (x for x in range(6) if x and False), + (y for x in range(6) if False and (y := x)), # noqa: F821 + ([x] for x in range(4) if False), + ((x, y) for x in range(4) for y in range(3) if False), + (x for x in range(6) if False if x > 1), + ({x for x in range(3)} for _ in range(2) if False), + ], +) +def test_unreachable_comprehension_body(genexpr): + """A comprehension whose body the compiler proved unreachable yields nothing.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + assert materialize(compile_and_eval(ast_node)) == [] + + +def test_outermost_iterable_strict_zip_stays_strict(): + """A ragged strict zip raises; the reconstruction must raise there too.""" + genexpr = (a + b for a, b in zip([1, 2, 3], [4, 5], strict=True)) + reconstructed = disassemble(genexpr) + assert ast.unparse(reconstructed) == ( + "(a + b for a, b in zip([1, 2, 3], [4, 5], strict=True))" + ) + + with pytest.raises(ValueError): + materialize(compile_and_eval(reconstructed)) + with pytest.raises(ValueError): + materialize(genexpr) + + +def test_outermost_iterable_lax_zip_stays_lax(): + """A lax zip stops at the shortest iterable and must not become strict.""" + genexpr = (a + b for a, b in zip([1, 2, 3], [4, 5])) + reconstructed = disassemble(genexpr) + assert ast.unparse(reconstructed) == "(a + b for a, b in zip([1, 2, 3], [4, 5]))" + assert materialize(compile_and_eval(reconstructed)) == [5, 7] + assert materialize(genexpr) == [5, 7] + + +def test_outermost_iterable_partially_consumed_strict_zip(): + """Strictness survives even once the zip has been partly consumed.""" + zipped = zip([1, 2, 3], [4, 5], strict=True) + next(zipped) + + genexpr = (a + b for a, b in zipped) + reconstructed = disassemble(genexpr) # must precede consuming `genexpr` + assert "strict=True" in ast.unparse(reconstructed) + with pytest.raises(ValueError): + materialize(compile_and_eval(reconstructed)) + + +@pytest.mark.parametrize( + "genexpr", + [ + # "dict_item" was once an internal marker in the first slot of a tuple, + # which is a string a user's data is perfectly entitled to hold. + (x for x in (("dict_item", 1, 2),)), + (x for x in [("dict_item", "key", "value")]), + (x for x in (("dict_item",), ("dict_item", 1), ("dict_item", 1, 2, 3))), + (("dict_item", x) for x in range(2)), + (("dict_item", x, x + 1) for x in range(2)), + (x for x in {("dict_item", 1, 2): "v"}.items()), + ], +) +def test_tuples_starting_with_dict_item(genexpr): + """No element of a user's tuple is an internal marker to be stripped.""" + assert_ast_equivalent(genexpr, disassemble(genexpr)) + + +def test_outermost_iterable_partially_consumed_adaptor(): + """A consumed prefix is reflected in the adaptor's inner iterators.""" + zipped = zip([1, 2, 3], [4, 5, 6]) + next(zipped) + + genexpr = (a + b for a, b in zipped) + reconstructed = disassemble(genexpr) # must precede consuming `genexpr` + assert materialize(genexpr) == [7, 9] + assert materialize(compile_and_eval(reconstructed)) == [7, 9] + + +def test_outermost_iterable_partially_consumed(): + """Only the *unconsumed* remainder of the outermost iterator belongs in the AST.""" + iterator = iter([10, 20, 30, 40]) + next(iterator) + next(iterator) + + genexpr = (x + 1 for x in iterator) + assert ast.unparse(disassemble(genexpr)) == "(x + 1 for x in [30, 40])" + assert materialize(genexpr) == [31, 41] + + +def test_outermost_iterable_partially_consumed_str(): + iterator = iter("hello") + next(iterator) + + genexpr = (c for c in iterator) + assert ast.unparse(disassemble(genexpr)) == "(c for c in 'ello')" + assert materialize(genexpr) == ["e", "l", "l", "o"] + + +# ============================================================================ +# BINARY OPERATORS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # Bitwise operators, which BINARY_OP folds in with the arithmetic ones + (x & 3 for x in range(8)), + (x | 3 for x in range(8)), + (x ^ 3 for x in range(8)), + (x << 2 for x in range(4)), + (x >> 1 for x in range(8)), + (~x & 7 for x in range(8)), + # Mixed precedence across the whole operator table + (x & 1 | x >> 2 ^ 3 for x in range(8)), + ((x | 1) & (x ^ 2) for x in range(8)), + (x + 1 & x - 1 for x in range(8)), + (x * 2 % 5 // 2 for x in range(8)), + (x**2 - x // 2 + x % 3 for x in range(1, 8)), + # Operators on non-numeric operands + (s + "!" for s in ["a", "b"]), + (s * 2 for s in ["a", "b"]), + (t + (9,) for t in [(1,), (2,)]), + (frozenset({x}) | frozenset({9}) for x in range(3)), + (frozenset({x, 1}) & frozenset({1}) for x in range(3)), + (frozenset({x, 1}) ^ frozenset({1}) for x in range(3)), + ({"a": x} | {"b": 0} for x in range(3)), + ], +) +def test_binary_operators(genexpr): + """Test reconstruction of the full BINARY_OP table.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +def test_matmul_operator(): + """BINARY_OP argument 4 is `@`, which no built-in type implements. + + The comprehension is disassembled but never evaluated, so this checks the + reconstructed source rather than the reconstructed values. + """ + genexpr = (a @ b for a in [1, 2]) # noqa: F821 + assert ast.unparse(disassemble(genexpr)) == "(a @ b for a in (1, 2))" + + +# ============================================================================ +# KEYWORD ARGUMENTS AT CALL SITES +# +# Python 3.12 compiles these as KW_NAMES followed by CALL; 3.13 replaced the +# pair with a single CALL_KW instruction. +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr,globals_dict", + [ + ((dict(a=x) for x in range(3)), {}), + ((dict(a=x, b=x * 2) for x in range(3)), {}), + ((dict(a=x if x > 0 else 0, b=x * 2 if x < 5 else x) for x in range(8)), {}), + # Mixed positional and keyword arguments + ((sorted([3, x], reverse=True) for x in range(3)), {}), + ((sorted([x, 1], key=abs) for x in range(3)), {}), + ((sorted([x, 1], key=abs, reverse=True) for x in range(3)), {}), + ((int(str(x), base=8) for x in range(8)), {}), + ((round(x / 3, ndigits=2) for x in range(5)), {}), + # Keyword arguments on a method call + (("a,b".split(sep=",") for x in range(2)), {}), + (("a-b".replace("-", "+") for x in range(2)), {}), + # Nested calls, each with keywords + ((dict(a=dict(b=x)) for x in range(3)), {}), + ((sorted(sorted([x, 1]), reverse=True) for x in range(3)), {}), + # Keyword arguments to a user-supplied callable + ((f(x, scale=2) for x in range(3)), {"f": lambda v, scale=1: v * scale}), # type: ignore # noqa: F821 + ], +) +def test_keyword_arguments(genexpr, globals_dict): + """Test reconstruction of calls with keyword arguments.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node, globals_dict) + + +@pytest.mark.parametrize( + "genexpr", + [ + # Starred positional arguments + (max(*[x, 1]) for x in range(3)), + (max(1, *[x, 2]) for x in range(3)), + (max(*[x, 1], *[2, 3]) for x in range(3)), + (max(*(x, 1)) for x in range(3)), + (sum([x, 1], *[0]) for x in range(3)), + # Double-starred keyword arguments + (dict(**{"a": x}) for x in range(3)), + (dict(a=x, **{"b": 1}) for x in range(3)), + (dict(**{"a": x}, **{"b": 2}) for x in range(3)), + (sorted([x, 1], **{"reverse": True}) for x in range(3)), # type: ignore[call-overload] + # Both at once + (max(*[[x, 1]], **{"default": 0}) for x in range(3)), # type: ignore[call-overload] + # On a method call, where the callable comes with a `self` + ("-".join(*[[str(x), "z"]]) for x in range(3)), + # Unpacking a comprehension, and unpacking into a nested call + (max(*[y for y in range(x + 2)]) for x in range(3)), + (max(*[abs(y) for y in range(-x - 1, 1)]) for x in range(3)), + (dict(**{str(k): k for k in range(x + 1)}) for x in range(3)), + ], +) +def test_star_argument_calls(genexpr): + """Test reconstruction of `*args`/`**kwargs` call sites (CALL_FUNCTION_EX). + + The arguments arrive already collected into a sequence and a mapping, so the + reconstruction spells every argument as unpacked; that evaluates identically + even where the source passed some of them plainly. + """ + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# LAMBDAS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + ((lambda y: y * 2)(x) for x in range(5)), + ((lambda y, z: y + z)(x, x) for x in range(5)), + ((lambda y: y)(x) for x in range(5)), + # Closing over the loop variable, and over an enclosing comprehension + ((lambda y: y + x)(x) for x in range(5)), + ((lambda: x)() for x in range(5)), + (((lambda y: lambda z: z + y)(x))(1) for x in range(5)), + # Lambdas whose body is itself a comprehension + ((lambda: (y for y in range(x)))() for x in range(3)), + ((lambda: [y for y in range(x)])() for x in range(3)), + ((lambda n: sum(y for y in range(n)))(x) for x in range(3)), + # Lambdas as arguments to other calls + (sorted([x, 1], key=lambda v: -v) for x in range(3)), + (list(map(lambda v: v * 2, [x, 1])) for x in range(3)), + # Conditional expressions inside a lambda body + ((lambda y: y if y % 2 else -y)(x) for x in range(5)), + ((lambda y: (y if y > 1 else 0) + 1)(x) for x in range(5)), + ], +) +def test_lambdas(genexpr): + """Test reconstruction of lambdas appearing inside comprehensions.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +@pytest.mark.parametrize( + "genexpr", + [ + # Positional defaults, which attach to the *trailing* parameters + ((lambda y, z=2: y * z)(x) for x in range(3)), # type: ignore[assignment] + ((lambda y=1: y)() for x in range(3)), # type: ignore[assignment] + ((lambda y, z=2, w=3: y * z * w)(x) for x in range(3)), # type: ignore[assignment] + ((lambda y, z=2: y * z)(x, 5) for x in range(3)), + # Keyword-only parameters, with and without defaults + ((lambda y, *, z=1: y + z)(x) for x in range(3)), # type: ignore[assignment] + ((lambda y, *, z=1, w=2: y + z + w)(x) for x in range(3)), # type: ignore[assignment] + ((lambda y, *, z: y + z)(x, z=4) for x in range(3)), + # Positional-only parameters + ((lambda y, /, z=2: y * z)(x) for x in range(3)), # type: ignore[assignment] + # *args and **kwargs + ((lambda *a: sum(a))(x, x) for x in range(3)), + ((lambda **k: sum(k.values()))(a=x) for x in range(3)), + ((lambda *a, **k: len(a) + len(k))(x, b=1) for x in range(3)), + ((lambda y, *a: y + len(a))(x, 1, 2) for x in range(3)), + ( + (lambda y, *a, z=3, **k: y + len(a) + z + len(k))(x, 1, w=2) + for x in range(3) + ), + # Defaults that are themselves non-trivial expressions + ((lambda y, z=(1, 2): y + len(z))(x) for x in range(3)), # type: ignore[assignment] + ((lambda y, z=[1]: y + len(z))(x) for x in range(3)), # type: ignore[assignment] + ], +) +def test_lambda_default_and_variadic_arguments(genexpr): + """Test reconstruction of lambda defaults and variadic parameters.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +@pytest.mark.parametrize( + "genexpr", + [ + # A lambda reached as a live object -- through the outermost iterable -- + # rather than built inside the comprehension. Its defaults are on the + # function, not in the code object it was compiled to. + (fn() for fn in [lambda value=1: value]), + (fn(2) for fn in [lambda a, b=10: a * b]), + (fn(2) for fn in (lambda a, b=10: a * b,)), + (fn(1) for fn in [lambda a, *, k=5: a + k]), + (fn(1) for fn in [lambda a, /, b=2, *, k=5: a * b + k]), + (fn() for fn in [lambda x=(1, 2): sum(x)]), + (fn() for fn in [lambda x=[1, 2]: len(x)]), + (fn(1, 2) for fn in [lambda a, b=0, *rest, k=5, **kw: a + b + k + len(rest)]), + # ... several of them, and one with no defaults alongside + (fn() for fn in [lambda: 0, lambda v=1: v, lambda v=2: v]), # type: ignore[misc] + # ... and one passed through map rather than iterated directly + (fn(3) for fn in map(lambda f: f, [lambda a, b=4: a * b])), + ], +) +def test_lambda_object_defaults(genexpr): + """A lambda arriving as a live object keeps the defaults attached to it.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# ASSIGNMENT EXPRESSIONS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # A walrus binds in the *enclosing* scope, so at module level it + # compiles to COPY + STORE_GLOBAL rather than to a local. + (y for x in range(4) if (y := x * 2) > 1), # noqa: F821 + ((y := x) + 1 for x in range(3)), # noqa: F821 + (y * y for x in range(4) if (y := x + 1) > 2), # noqa: F821 + ((y := x) if x > 1 else -1 for x in range(4)), # noqa: F821 + # Bound in one clause and read in a later one + ((y, z) for x in range(3) if (y := x + 1) for z in range(y)), # noqa: F821 + # Inside a nested comprehension, and inside a lambda + ([(z := w) + z for w in range(x)] for x in range(4)), # noqa: F821 + ((lambda n: [(z := w) + z for w in range(n)])(x) for x in range(4)), # noqa: F821 + # Combined with a short-circuiting filter. The `or` must not re-evaluate + # the assignment, which is why the disjunction is absorbed. + (y for x in range(6) if (y := x * 2) > 6 or y == 0), # noqa: F821 + ], +) +def test_assignment_expressions(genexpr): + """A walrus in a comprehension binds in the *enclosing* scope (STORE_GLOBAL here).""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# COMPREHENSIONS NESTED IN EACH SYNTACTIC POSITION +# +# A comprehension can appear in the element, in the iterable, and inside a +# filter, and each of the four comprehension kinds can nest inside any other. +# The filter position is the interesting one: filters are reconstructed from +# control flow, so a comprehension inside a filter has to survive being treated +# as part of a boolean condition. +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # ... in filter position + (x for x in range(5) if any(y > 2 for y in range(x))), + (x for x in range(5) if all(y < 3 for y in range(x))), + (x for x in range(5) if [y for y in range(x)]), + (x for x in range(5) if {y for y in range(x)}), + (x for x in range(5) if {y: y for y in range(x)}), + (x for x in range(6) if len([y for y in range(x) if y % 2]) > 1), + (x for x in range(5) if any(y for y in range(x) if y % 2)), + ( + x + for x in range(6) + if all(y < x for y in range(2)) and any(z > 1 for z in range(x)) + ), # noqa: E501 # fmt: skip + # ... in filter position, inside a short-circuiting condition + (x for x in range(6) if sum(y for y in range(x)) > 3 or x == 0), + (x for x in range(6) if x == 0 or any(y > 2 for y in range(x))), + (x for x in range(6) if x == 0 or len([y for y in range(x)]) > 2), + ( + x + for x in range(6) + if any(y > 1 for y in range(x)) or all(z < 2 for z in range(x)) + ), # noqa: E501 # fmt: skip + # ... in iterable position + (x for x in [y for y in range(5) if y % 2]), + (x for x in {y for y in range(5) if y % 2}), + (x for x in {y: y for y in range(3)}), + (x for x in (y for y in range(5) if y > 1 or y == 0)), + (x for x in [y for y in [z for z in range(4)] if y % 2]), + (x for x in [y for y in range(4)] if x > 1 or x == 0), + ((a, b) for a in range(3) for b in [c for c in range(a)]), + # ... in element position + ([y for y in range(x) if y % 2 or y == 0] for x in range(4)), + ({y for y in range(x) if y > 1} for x in range(5)), + ({y: [z for z in range(y)] for y in range(x)} for x in range(4)), + (sum(y for y in range(x) if y % 2) for x in range(5)), + ([(z for z in range(y)) for y in range(x)] for x in range(3)), + # ... in several positions at once, with different kinds + ( + [y for y in {z for z in range(x)}] + for x in range(4) + if any(w > 1 for w in range(x)) + ), # noqa: E501 # fmt: skip + ({k: [v for v in range(k)] for k in {j for j in range(x)}} for x in range(4)), + ([y for y in range(x) if y or y == 0] for x in range(5) if x > 1 or x == 0), + ( + (y for y in range(x) if y % 2 or y == 0) + for x in (z for z in range(4)) + if x < 3 or x == 3 + ), # noqa: E501 # fmt: skip + # ... nesting the four kinds inside one another + ({k: {v for v in range(k)} for k in [j for j in range(x)]} for x in range(4)), + ([{y: y} for y in range(x)] for x in range(4)), + ({(y, y * 2) for y in range(x)} for x in range(4)), + ({y: y for y in range(x) if y % 2 or y == 0} for x in range(5)), + ({y for y in range(x) if y % 2 or y == 0} for x in range(5)), + ([y for y in range(x) if 1 < y < 4] for x in range(6)), + ], +) +def test_comprehensions_in_every_position(genexpr): + """Test comprehensions nested in the element, the iterable and the filter.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +@pytest.mark.parametrize( + "genexpr", + [ + # A lambda body is a separate code object with no loop of its own, so + # every branch inside one is a conditional expression rather than a + # filter -- including the branches of a comprehension nested in it. + ((lambda n: [y for y in range(n) if y % 2])(x) for x in range(4)), + ((lambda n: {y for y in range(n)})(x) for x in range(4)), + ((lambda n: {y: y**2 for y in range(n)})(x) for x in range(4)), + ((lambda n: sum(y for y in range(n)))(x) for x in range(4)), + ((lambda n: [y for y in range(n) if y > 1 or y == 0])(x) for x in range(5)), + ((lambda n: (y for y in range(n) if y % 2 or y == 0))(x) for x in range(4)), + ((lambda n: [y if y > 1 else -y for y in range(n)])(x) for x in range(4)), + ((lambda n: [y for y in range(n) if 1 < y < 3])(x) for x in range(5)), + # Lambdas nested in lambdas, and lambdas inside the comprehension body + ((lambda n: (lambda m: [y for y in range(m)])(n))(x) for x in range(3)), + ([(lambda v: v * 2)(y) for y in range(x)] for x in range(4)), + (list(map(lambda n: [y for y in range(n)], range(x))) for x in range(3)), + ([(lambda v: v if v > 1 else -v)(y) for y in range(x)] for x in range(4)), + ], +) +def test_comprehensions_inside_lambdas(genexpr): + """Test comprehensions nested inside lambda bodies.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# CLOSURES +# +# A comprehension written inside a function reads the function's locals out of +# closure cells, not out of globals. Those cells belong to a scope the +# reconstruction does not reproduce, so leaving a free variable as a bare name +# would silently turn it into a global lookup -- picking up a different value, +# or none at all. The captured value is written into the tree instead. Each +# test below evaluates the reconstruction in a namespace that binds the same +# name to something else, so a lookup that leaked out would be caught. +# ============================================================================ + + +def test_closure_shadowed_by_global(): + def make(): + value = 1 + return (value for _ in range(1)) + + genexpr = make() + reconstructed = disassemble(genexpr) + assert ast.unparse(reconstructed) == "(1 for _ in range(0, 1, 1))" + assert materialize(genexpr) == [1] + assert materialize(compile_and_eval(reconstructed, {"value": 2})) == [1] + + +@pytest.mark.parametrize( + "make,shadow,expected", + [ + # In the element expression, the filter, and an inner iterable + (lambda: (lambda n: (x * n for x in range(4)))(3), {"n": 100}, [0, 3, 6, 9]), + ( + lambda: (lambda t: (x for x in range(6) if x > t))(3), + {"t": -1}, + [4, 5], + ), + ( + lambda: (lambda k: (y for x in range(2) for y in range(k)))(2), + {"k": 5}, + [0, 1, 0, 1], + ), + # A captured container, indexed and iterated + ( + lambda: (lambda d: (d[i] for i in range(2)))([10, 20]), + {"d": [0, 0]}, + [10, 20], + ), + ( + lambda: (lambda d: (v for v in d))({"a": 1, "b": 2}), + {"d": {}}, + ["a", "b"], + ), + # Captured by a lambda nested inside the comprehension + ( + lambda: (lambda n: ((lambda y: y * n)(x) for x in range(4)))(3), + {"n": 100}, + [0, 3, 6, 9], + ), + ( + lambda: (lambda n: ((lambda: (lambda: n)())() for _ in range(2)))(7), + {"n": 0}, + [7, 7], + ), + # Captured by a comprehension nested inside the comprehension + ( + lambda: (lambda n: ([y * n for y in range(x)] for x in range(3)))(2), + {"n": 100}, + [[], [0], [0, 2]], + ), + ( + lambda: (lambda n: (sum(y for y in range(x) if y < n) for x in range(4)))( + 2 + ), + {"n": 100}, + [0, 0, 1, 1], + ), + # Two free variables at once + ( + lambda: (lambda a, b: (x * a + b for x in range(3)))(2, 1), + {"a": 0, "b": 0}, + [1, 3, 5], + ), + ], +) +def test_closure_free_variables(make, shadow, expected): + genexpr = make() + reconstructed = disassemble(genexpr) + assert materialize(genexpr) == expected + assert materialize(compile_and_eval(reconstructed, dict(shadow))) == expected + # The name must not survive anywhere in the tree, or the shadowing binding + # above would have been the one that answered. + assert not ( + {node.id for node in ast.walk(reconstructed) if isinstance(node, ast.Name)} + & set(shadow) + ) + + +def test_closure_target_captured_by_nested_lambda_stays_a_name(): + """A cell this comprehension *creates* is bound in the reconstruction too.""" + genexpr = ((lambda: x)() for x in range(3)) + reconstructed = disassemble(genexpr) + assert materialize(genexpr) == [0, 1, 2] + assert materialize(compile_and_eval(reconstructed, {"x": 99})) == [0, 1, 2] + + +def test_closure_shadowed_by_a_nested_comprehension_target(): + """Only the free `n` is the captured one; the inner comprehension rebinds it.""" + + def make(): + n = 5 + return (n + sum(n for n in range(x)) for x in range(3)) + + genexpr = make() + reconstructed = disassemble(genexpr) + assert materialize(genexpr) == [5, 5, 6] + assert materialize(compile_and_eval(reconstructed, {"n": 100, "sum": sum})) == [ + 5, + 5, + 6, + ] + + +def test_closure_shadowed_by_a_nested_lambda_parameter(): + """Only the free `n` is the captured one; the lambda's parameter is its own.""" + + def make(): + n = 5 + return ((lambda n: n * 2)(x) + n for x in range(3)) + + genexpr = make() + reconstructed = disassemble(genexpr) + assert materialize(genexpr) == [5, 7, 9] + assert materialize(compile_and_eval(reconstructed, {"n": 100})) == [5, 7, 9] + + +def test_closure_inside_a_live_lambda(): + """A lambda reached as an object closes over cells of its own.""" + + def make(): + n = 3 + return (fn(2) for fn in [lambda a, b=10: a * n + b]) + + genexpr = make() + reconstructed = disassemble(genexpr) + assert materialize(genexpr) == [16] + assert materialize(compile_and_eval(reconstructed, {"n": 100})) == [16] + + +def test_closure_value_that_cannot_be_represented(): + """A capture with no AST spelling is refused rather than quietly dropped.""" + + class Opaque: + pass + + def make(): + obj = Opaque() + return (obj for _ in range(1)) + + with pytest.raises(TypeError, match="captured in free variable 'obj'"): + disassemble(make()) + + +def test_closure_captured_iterator_is_refused(): + """An iterator's remaining elements are not the iterator, so it is refused.""" + + def make(): + flags = iter([True, False, True, False]) + return (x for x in range(4) if next(flags)) + + with pytest.raises(TypeError, match="captured in free variable 'flags'"): + disassemble(make()) + + +# ============================================================================ +# STATEFUL EXPRESSIONS +# +# What comes back is syntax, so evaluating it runs the comprehension a second +# time. A filter or element expression that depends on state that has since +# moved on answers differently then -- faithfully reconstructed, but no longer +# in agreement with the generator it came from. +# ============================================================================ + + +_FLAGS = iter([True, False, True, False, False, False, True, False]) + + +def test_stateful_filter_is_re_evaluated(): + genexpr = (x for x in range(4) if next(_FLAGS)) + reconstructed = disassemble(genexpr) + assert ast.unparse(reconstructed) == "(x for x in range(0, 4, 1) if next(_FLAGS))" + + # The reconstruction is the same comprehension, but `_FLAGS` has advanced by + # the time it runs, so the two do not agree element for element. + assert materialize(genexpr) == [0, 2] + assert materialize( + compile_and_eval(reconstructed, {"_FLAGS": _FLAGS, "next": next}) + ) == [2] + + +@pytest.mark.parametrize( + "genexpr", + [ + # Short-circuiting conditions nested in one another + (x for x in range(20) if (x > 2 or x < 1) and (x < 10 or x > 15)), + (x for x in range(20) if ((x > 2 and x < 5) or (x > 10 and x < 15)) or x == 0), + (x for x in range(30) if not (x % 2 == 0 or x % 3 == 0)), + (x for x in range(30) if not (not (x > 5) or not (x < 20))), + (x for x in range(40) if (x > 5 and x < 35) and (x % 3 == 0 or x % 5 == 0)), + ( + x + for x in range(40) + if (x < 5 and x % 2 == 0) or (10 < x < 15) or (x > 35 and x % 3 == 0) + ), # noqa: E501 # fmt: skip + # Short-circuiting conditions spanning several generators + ( + (x, y) + for x in range(6) + if x < 2 or x > 4 + for y in range(6) + if y < 1 or y > 4 + ), + ( + (x, y) + for x in range(5) + if x % 2 == 0 or x == 1 + for y in range(x) + if y > 0 and y < 3 + ), # noqa: E501 # fmt: skip + # Conditional expressions and filters that are both lazy + ((x if (x > 2 or x < 1) else -x) for x in range(10) if x % 2 == 0 or x == 1), + (x for x in range(20) if (x if x > 5 else not x) or x == 3), + ( + (x if x > 5 or x < 2 else (0 if x % 2 == 0 or x == 3 else 1)) + for x in range(12) + ), # noqa: E501 # fmt: skip + # Chained comparisons combined with lazy operators + (x for x in range(30) if 5 < x < 15 or 20 < x < 25), + (x for x in range(30) if 5 < x < 15 and (x % 2 == 0 or x % 3 == 0)), + ((x if 5 < x < 15 else 0) for x in range(20) if 2 < x < 18), + # Lazy conditions inside a nested comprehension, and around it + ([y for y in range(x) if y > 1 or y == 0] for x in range(5) if x > 2 or x == 0), + ((y for y in range(x) if y % 2 or y == 0) for x in range(4) if x < 3 or x == 3), + ], +) +def test_nested_lazy_conditions(genexpr): + """Test short-circuiting conditions nested inside one another.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# STRUCTURAL STRESS CASES +# ============================================================================ + +# These two must stay on one line: on Python 3.12 `dis` mis-reports jumps for +# multiline comprehensions, which test_multiline_comprehensions covers directly. +_STRESS_MANY_FILTERS = ((x, y) for x in range(10) if x % 2 == 0 if x > 2 for y in range(10) if y % 3 == 0 if y < x) # fmt: skip +_STRESS_NESTED_TERNARY = ([y if y > 1 else -y for y in range(x)] for x in range(4) if (x if x % 2 == 1 else x % 2 == 0)) # fmt: skip + + +@pytest.mark.parametrize( + "genexpr", + [ + # Deep loop nesting + ( + a + b + c + d + e + for a in range(2) + for b in range(2) + for c in range(2) + for d in range(2) + for e in range(2) + ), + ( + (a, b, c, d) + for a in range(2) + for b in range(a + 1) + for c in range(b + 1) + for d in range(c + 1) + ), + # Many filters spread over many loops. Kept on one line: on Python 3.12 + # `dis` mis-reports jumps for multiline comprehensions, which is covered + # separately by test_multiline_comprehensions below. + (x for x in range(50) if x > 5 if x < 40 if x % 2 == 0 if x % 3 == 0), + _STRESS_MANY_FILTERS, + # Deep comprehension nesting + (((z for z in range(y)) for y in range(x)) for x in range(3)), + ([[z for z in range(y)] for y in range(x)] for x in range(3)), + ({y: [z for z in range(y)] for y in range(x)} for x in range(3)), + # Structured literals in the element position + ({x, x + 1} for x in range(3)), + ({x: x + 1} for x in range(3)), + (((x, x), x) for x in range(3)), + ([x, [x, [x]]] for x in range(3)), + ({"k": [x, {"j": (x,)}]} for x in range(3)), + # Ternaries interleaved with nesting + _STRESS_NESTED_TERNARY, + (((y if y else -1) for y in range(x)) for x in range(3)), + ], +) +def test_structural_stress(genexpr): + """Test reconstruction of deeply nested and heavily filtered comprehensions.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# MULTILINE COMPREHENSIONS +# +# On Python 3.12 `dis` reports a different jump layout for a filter whose source +# spans several lines. The reconstruction used to come out negated, because the +# filter/conditional distinction was drawn from the *local* instruction order. +# Classifying branches from the control-flow graph instead is insensitive to +# that, so these now reconstruct identically on 3.12 and 3.13+. +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + ( + x + for x in range(5) # comment to avoid reformatting + if x > 1 + ), + ( + x + for x in range(10) # comment to avoid reformatting + if x % 2 == 0 + if x > 2 + ), + ( + (x, y) + for x in range(10) + if x % 2 == 0 + if x > 2 + for y in range(10) + if y % 3 == 0 + if y < x + ), + ( + [y if y > 1 else -y for y in range(x)] + for x in range(4) + if (x if x % 2 == 1 else x % 2 == 0) + ), + ], +) +def test_multiline_comprehensions(genexpr): + """Filters in multiline comprehensions are mis-disassembled on Python 3.12.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +def test_multiline_comprehensions_same_on_one_line(): + """The same expressions reconstruct correctly when written on one line.""" + one_line = (x for x in range(10) if x % 2 == 0 if x > 2) + assert_ast_equivalent(one_line, disassemble(one_line)) + + +# ============================================================================ +# HELPER FUNCTION TESTS +# ============================================================================ + + +@pytest.mark.parametrize( + "value,expected_str", + [ + # AST nodes should be returned as-is + (ast.Name(id="x", ctx=ast.Load()), "x"), + (ast.Constant(value=42), "42"), + (ast.List(elts=[], ctx=ast.Load()), "[]"), + ( + ast.BinOp( + left=ast.Constant(value=1), op=ast.Add(), right=ast.Constant(value=2) + ), + "1 + 2", + ), + # Constants should become ast.Constant nodes + (42, "42"), + (3.14, "3.14"), + (-42, "-42"), + (-3.14, "-3.14"), + ("hello", "'hello'"), + ("", "''"), + (b"bytes", "b'bytes'"), + (b"", "b''"), + (True, "True"), + (False, "False"), + (None, "None"), + # Complex numbers + (1 + 2j, "(1+2j)"), + (0 + 1j, "1j"), + (3 + 0j, "(3+0j)"), + (-1 - 2j, "(-1-2j)"), + # Tuples should become ast.Tuple nodes + ((), "()"), + ((1,), "(1,)"), + ((1, 2), "(1, 2)"), + (("a", "b", "c"), "('a', 'b', 'c')"), + # A tuple is a tuple: no element of one is a marker to be stripped + (("dict_item", "key", "value"), "('dict_item', 'key', 'value')"), + (("dict_item", 42, "answer"), "('dict_item', 42, 'answer')"), + # Nested tuples + ((1, (2, 3)), "(1, (2, 3))"), + (((1, 2), (3, 4)), "((1, 2), (3, 4))"), + ((1, 2, (3, (4, 5))), "(1, 2, (3, (4, 5)))"), + # Lists should become ast.List nodes + ([1, 2, 3], "[1, 2, 3]"), + (["hello", "world"], "['hello', 'world']"), + ([True, False, None], "[True, False, None]"), + # Nested lists + ([[1, 2], [3, 4]], "[[1, 2], [3, 4]]"), + ([1, [2, [3, 4]], 5], "[1, [2, [3, 4]], 5]"), + # Mixed nested structures + ([(1, 2), (3, 4)], "[(1, 2), (3, 4)]"), + (([1, 2], [3, 4]), "([1, 2], [3, 4])"), + # Dicts should become ast.Dict nodes + ({"a": 1}, "{'a': 1}"), + ({"x": 10, "y": 20}, "{'x': 10, 'y': 20}"), + ({1: "one", 2: "two"}, "{1: 'one', 2: 'two'}"), + # Nested dicts + ({"a": {"b": 1}}, "{'a': {'b': 1}}"), + ( + {"nums": [1, 2, 3], "strs": ["a", "b"]}, + "{'nums': [1, 2, 3], 'strs': ['a', 'b']}", + ), + # Range objects + (range(5), "range(0, 5, 1)"), + (range(1, 10), "range(1, 10, 1)"), + (range(0, 10, 2), "range(0, 10, 2)"), + (range(10, 0, -1), "range(10, 0, -1)"), + (range(-5, 5), "range(-5, 5, 1)"), + # Empty collections + ([], "[]"), + ((), "()"), + ({}, "{}"), + # Complex nested structures + ([1, [2, 3], 4], "[1, [2, 3], 4]"), + ({"a": [1, 2], "b": {"c": 3}}, "{'a': [1, 2], 'b': {'c': 3}}"), + ([(1, {"a": [2, 3]}), ({"b": 4}, 5)], "[(1, {'a': [2, 3]}), ({'b': 4}, 5)]"), + # Edge cases with special values + ([None, True, False, 0, ""], "[None, True, False, 0, '']"), + ( + {"": "empty", None: "none", 0: "zero"}, + "{'': 'empty', None: 'none', 0: 'zero'}", + ), + # Large numbers + (999999999999999999999, "999999999999999999999"), + (1.7976931348623157e308, "1.7976931348623157e+308"), # Close to float max + # Sets - note unparse equivalence may fail for unordered collections + ({1, 2, 3}, "{1, 2, 3}"), + ], +) +def test_ensure_ast(value, expected_str): + """Test that ensure_ast correctly converts various values to AST nodes.""" + + result = ensure_ast(value) + + # Compare the unparsed strings + result_str = ast.unparse(result) + assert result_str == expected_str, ( + f"ensure_ast({repr(value)}) produced '{result_str}', expected '{expected_str}'" + ) + + +def test_error_handling(): + """Test that appropriate errors are raised for unsupported cases.""" + # Test with non-generator input + with pytest.raises(ValueError): + disassemble([1, 2, 3]) # Not a generator + + # Test with consumed generator + gen = (x for x in range(5)) + list(gen) # Consume it + with pytest.raises(ValueError): + disassemble(gen) + + # Test with a generator that has been started but not consumed + gen = (x for x in range(5)) + next(gen) + with pytest.raises(ValueError): + disassemble(gen) + + +def test_comp_lambda_copy(): + """Test that CompLambda is compatible with copy.copy and copy.deepcopy.""" + # Create a test generator expression AST + genexpr_ast = ast.GeneratorExp( + elt=ast.Name(id="x", ctx=ast.Load()), + generators=[ + ast.comprehension( + target=ast.Name(id="x", ctx=ast.Store()), + iter=DummyIterName(), + ifs=[], + is_async=0, + ) + ], + ) + + # Create a CompLambda instance + comp_lambda = CompLambda(genexpr_ast) + + # Test copy.copy + copied = copy.copy(comp_lambda) + assert isinstance(copied, CompLambda) + assert ast.unparse(copied.body) == ast.unparse(comp_lambda.body) + assert copied.body is comp_lambda.body # Shallow copy shares the body + + # Test copy.deepcopy + deep_copied = copy.deepcopy(comp_lambda) + assert isinstance(deep_copied, CompLambda) + assert ast.unparse(deep_copied.body) == ast.unparse(comp_lambda.body) + assert deep_copied.body is not comp_lambda.body # Deep copy creates new body + + # Test that deep copied version works the same way + iterator = ast.Call( + func=ast.Name(id="range", ctx=ast.Load()), + args=[ast.Constant(value=5)], + keywords=[], + ) + + original_result = comp_lambda.inline(iterator) + deep_copied_result = deep_copied.inline(iterator) + + assert ast.unparse(original_result) == ast.unparse(deep_copied_result) + assert type(original_result) == type(deep_copied_result) + + +# ============================================================================ +# AST TRANSFORMER TESTS +# ============================================================================ From c074f46d7b35b2efe2a1cd406c64a478a8ed48e0 Mon Sep 17 00:00:00 2001 From: Eli Date: Tue, 28 Jul 2026 12:00:11 -0400 Subject: [PATCH 15/16] Add ReduceGroundCartesianProduct Replaces a cartesian-product stream with one stream per plate assignment the body actually subscripts. Where the body folds over the plates uniformly, ReduceDistributeCartesianProduct inverts the reduction instead, which is cheaper and leaves the plate fold intact. This rule is the fallback for bodies that admit no per-plate factorization: a chain `X[t], X[t+1]` couples adjacent plate indices, so inversion does not apply. Grounding costs one variable per assignment instead of enumerating the |D|^|P| rows, after which the result is an ordinary variable-elimination problem that `Factor` solves. `ReducePartial` gains an `unrolled` filter so grounding can expand only the plate folds it is grounding over. Co-Authored-By: Claude Opus 5 (1M context) --- effectful/ops/monoid.py | 208 +++++++++++++++++++++++++++++++++++++++ tests/test_ops_monoid.py | 73 ++++++++++++++ 2 files changed, 281 insertions(+) diff --git a/effectful/ops/monoid.py b/effectful/ops/monoid.py index 4830133d3..b0b9256ae 100644 --- a/effectful/ops/monoid.py +++ b/effectful/ops/monoid.py @@ -20,6 +20,7 @@ as_dict, defdata, deffn, + defop, implements, ite, range_, @@ -571,12 +572,19 @@ def _(self, monoid, body, streams): class ReducePartial(ObjectInterpretation): + unrolled: collections.abc.Set[Operation] | None + + def __init__(self, unrolled: collections.abc.Set[Operation] | None = None): + self.unrolled = unrolled + @implements(Monoid.reduce) def _(self, monoid, body, streams): if not streams: return fwd() for stream_key, stream_body, streams_tail in outer_stream(streams): + if self.unrolled is not None and stream_key not in self.unrolled: + continue if isinstance(stream_body, Term): continue stream_values_iter = iter(stream_body) @@ -1040,6 +1048,205 @@ def _to_body(mapping, key): return fwd() +class _GroundRow(collections.abc.Mapping): + """A cartesian-product row that mints a variable per index it is asked for. + + Standing this in for a row variable passes a concrete key straight + to :meth:`__getitem__` here, which answers with the variable for that plate + assignment -- minting it, and the stream it ranges over, on first sight. + Asking twice for the same assignment gives the same variable, which is what + ties the factors of a chain together. + + A key that is still symbolic never reaches here at all: the subscript stays + a term until the fold supplying the index is expanded, after which the same + traversal passes over it again with the key concrete. + """ + + variables: dict[tuple, Operation[[], typing.Any]] + streams: dict[Operation[[], typing.Any], collections.abc.Iterable] + sealed: bool + + def __init__( + self, + cprod_body: collections.abc.Iterable, + cprod_streams: dict[Operation, collections.abc.Iterable], + ): + """cprod_body is the body and cprod_streams is ths streams of a CartesianProduct.reduce expression""" + shape = self._row_shape(cprod_body) + assert shape is not None + (self._index, self._value), self._value_streams = shape + self._plates = cprod_streams + self.variables = {} + self.streams = {} + # A subscript whose index is not concrete yet holds the row until the + # fold supplying that index is expanded, so while grounding is under + # way the row does legitimately sit inside a term. + self.sealed = False + + def __getitem__(self, key: tuple) -> Term: + key = key if isinstance(key, tuple) else (key,) + if key not in self.variables: + fresh = defop(self._value.op) + self.streams[fresh] = self._value_streams[self._value.op] + self.variables[key] = fresh + return self.variables[key]() + + @staticmethod + def _row_shape(cprod_body) -> tuple[tuple[tuple, Term], Mapping] | None: + """The index and value of a cartesian product's rows, with the value's + domains, or ``None`` if the body is not a row stream. + + A row stream is a union of one-entry rows: for each value its domain takes, + a row mapping ``idx`` to that value. The value must be a bare call of its + domain variable, since a key's variable is that variable renamed. + """ + match cprod_body: + case Term( + Union.reduce, + ( + [Term(effectful.ops.syntax.as_dict, ((idx, value),), {})], + Mapping() as value_streams, + ), + {}, + ) if ( + isinstance(idx, Sequence) + and all(isinstance(i, Term) and not i.args for i in idx) + and isinstance(value, Term) + and not value.args + and value.op in value_streams + ): + return (tuple(idx), value), value_streams + case _: + return None + + def residual(self) -> Mapping | None: + """What is left of the row once the keys the body used are taken out. + + Grounding consumes the keys the body subscripts; the rest are still + assigned a value by every row, so they still multiply the number of + rows. Rather than accounting for them, they stay in the expression -- + the same cartesian product over the plate values no key consumed, which + :class:`ReduceDistributeCartesianProduct` also leaves behind when it + peels one plate off a product over several. + + ``None`` if what is left cannot be written as a cartesian product: the + complement of a set of keys is only a product of plate ranges when + there is one plate. + """ + consumed = { + plate: {key[position] for key in self.variables} + for position, plate in enumerate(i.op for i in self._index) + } + remaining = { + plate: [v for v in values if v not in consumed.get(plate, ())] + for plate, values in self._plates.items() + } + if not any(remaining.values()): + return {} + if len(self._plates) > 1: + return None + return {plate: values for (plate, values) in remaining.items() if values} + + def __iter__(self): + return iter(()) + + def __len__(self) -> int: + return 0 + + +@evaluate.register +def _(expr: _GroundRow): + if expr.sealed: + raise ValueError(f"{expr} should not have lived long enough to get here") + return expr + + +class ReduceGroundCartesianProduct(ObjectInterpretation): + """Replace a cartesian-product stream with one stream per plate assignment. + + reduce(M, body, {X: CartesianProduct.reduce(Union.reduce([as_dict((idx, w))], D), P)} ∪ S) + ═══════════════════════════════════════════════════════════════════════════ + reduce(M, body[X[a] := w_a], {w_a: D[idx := a]} ∪ S) + + where ``a`` ranges over the plate assignments the body actually subscripts. + + Where the body folds over the plates uniformly, + :class:`ReduceDistributeCartesianProduct` inverts the reduction instead, + which is cheaper and leaves the plate fold intact. This rule is the fallback + for bodies that admit no per-plate factorization -- a chain ``X[t], X[t+1]`` + couples adjacent plate indices, so inversion does not apply. Grounding it + costs one variable per assignment (``|P|`` variables over ``D``) instead of + enumerating the ``|D|^|P|`` rows, after which the resulting reduction over + ordinary range streams is an ordinary variable-elimination problem that + :class:`Factor` solves. + """ + + @staticmethod + def _plate_folds(body, monoid) -> collections.abc.Set[Operation]: + """The stream variables of the product folds nested in this reduction""" + folds: tuple[Term, ...] + if isinstance(body, Term) and _is_monoid_reduce(body.op): + inner_monoid, folds = body.op.__self__, (body,) + elif isinstance(body, Term) and _is_monoid_plus(body.op): + inner_monoid = body.op.__self__ + folds = tuple( + arg + for arg in body.args + if isinstance(arg, Term) and arg.op is inner_monoid.reduce + ) + else: + return frozenset() + + if not folds or not distributes_over(inner_monoid, monoid): + return frozenset() + return {var for fold in folds for var in typing.cast(Mapping, fold.args[1])} + + @implements(Monoid.reduce) + def reduce(self, monoid, body, streams): + plate_vars = self._plate_folds(body, monoid) + if not plate_vars: + return fwd() + + for var, stream in streams.items(): + if ( + isinstance(stream, Term) + and stream.op is CartesianProduct.reduce + and _GroundRow._row_shape(stream.args[0]) is not None + ): + row = _GroundRow(stream.args[0], stream.args[1]) + else: + continue + + rest = {k: v for (k, v) in streams.items() if k is not var} + with handler(ReducePartial(plate_vars)), handler({var: lambda: row}): + grounded, tail = evaluate((body, rest)) + + # Past this point a row in the result is a subscript whose index + # never became concrete, which is a half-ground term rather than a + # step of the rewrite. + row.sealed = True + try: + fvsof((grounded, tail)) + except ValueError: + continue + + # Whatever the body did not consume stays a cartesian product over + # the plate values no key took, still bound to the row variable, so + # its multiplicity is carried by the expression rather than counted + # here. Nothing is left when every key was consumed, which is the + # ordinary case. + plates = row.residual() + if plates is None or not row.variables: + # Nothing was consumed, so rebinding the row would reproduce + # the term this rule was given. + continue + residual = ( + {var: CartesianProduct.reduce(stream.args[0], plates)} if plates else {} + ) + return monoid.reduce(grounded, {**row.streams, **residual, **tail}) + return fwd() + + class ReduceUnion(ObjectInterpretation): @implements(Monoid.reduce) def reduce(self, monoid, body, streams): @@ -1752,6 +1959,7 @@ def extend(self, *intps: Interpretation) -> typing.Self: ReduceUnion(), ReduceSplit(), Factor(), + ReduceGroundCartesianProduct(), ReduceDistributeCartesianProduct(), ReduceWeightedStream(), ReduceMaskHoist(), diff --git a/tests/test_ops_monoid.py b/tests/test_ops_monoid.py index a2d2df2e3..8645c3e52 100644 --- a/tests/test_ops_monoid.py +++ b/tests/test_ops_monoid.py @@ -1,5 +1,6 @@ import functools import math +import random import sys import typing from collections.abc import Iterable, Mapping @@ -1359,3 +1360,75 @@ def test_reduce_unfactor_reduces(Sum, Product, backend: Backend): ) rhs = Sum.reduce(Product.plus(f(x()), g(y())), {x: X(), y: Y(), z: Z()}) backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceUnfactor()) + + +@pytest.mark.parametrize("T,K", [(10, 3)]) +def test_ground_cartesian_product_chain(T, K): + """A chain over a cartesian product grounds into a forward pass. + + ``ReduceDistributeCartesianProduct`` cannot invert this reduction: the body + subscripts the row at ``t`` and at ``t + 1``, so no single plate index + factors out. ``ReduceGroundCartesianProduct`` instead mints one variable + per plate assignment the body actually asks for, leaving an ordinary + variable-elimination problem whose elimination order is the forward + algorithm. + + Written out longhand; the generator-comprehension spelling + + .. code-block:: python + + Sum( + Product(phi()[t][ixs[t]][ixs[t + 1]] for t in range(T - 1)) + for ixs in CartesianProduct(range(K) for _ in range(T)) + ) + + desugars to exactly the reduction below, but the desugaring itself lands + separately. + + ``T`` is 10 rather than the 20 this scales to because normalization is + superlinear in the chain length until the ``evaluate`` fast path lands; at + 10 the whole test takes a few seconds, at 20 it takes over an hour. + """ + fs = [ + [[random.uniform(0, 1) for _ in range(K)] for _ in range(K)] for _ in range(T) + ] + + @Operation.define + def phi() -> list[list[list[float]]]: + raise NotHandled + + t = Operation.define(int, name="t") + plate = Operation.define(int, name="plate") + row_value = Operation.define(int, name="row_value") + ixs = Operation.define(Mapping[tuple, int], name="ixs") + + with handler(NormalizeIntp): + zf_normal = Sum.reduce( + Product.reduce( + phi()[t()][ixs()[(t(),)]][ixs()[(t() + 1,)]], + {t: range(T - 1)}, + ), + { + ixs: CartesianProduct.reduce( + Union.reduce( + [as_dict(((plate(),), row_value()))], {row_value: range(K)} + ), + {plate: range(T)}, + ) + }, + ) + + # Grounding is what does the work here: no other rule applies to this + # reduction, so without it normalization is the identity and the answer is + # only recovered by enumerating all ``K ** T`` rows during evaluation. + assert isinstance(zf_normal, Term) and zf_normal.op is Sum.reduce + assert list(zf_normal.args[1].values()) == [range(K)] + + with handler(EvaluateIntp), handler({phi: lambda: fs}): + zf_actual = evaluate(zf_normal) + + alpha = [1.0] * K + for t_ in range(T - 1): + alpha = [sum(alpha[i] * fs[t_][i][j] for i in range(K)) for j in range(K)] + zf_expected = sum(alpha) + assert isinstance(zf_actual, float) and math.isclose(zf_actual, zf_expected) From 7be9b05b6382c28cb72f7b11b934bdd0e6d1b3aa Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Wed, 29 Jul 2026 10:00:47 -0400 Subject: [PATCH 16/16] add more extensive tests --- tests/test_ops_monoid.py | 121 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 121 insertions(+) diff --git a/tests/test_ops_monoid.py b/tests/test_ops_monoid.py index 8645c3e52..4dffba967 100644 --- a/tests/test_ops_monoid.py +++ b/tests/test_ops_monoid.py @@ -37,6 +37,7 @@ ReduceEmpty, ReduceEqualityMaskRange, ReduceFusion, + ReduceGroundCartesianProduct, ReduceMaskHoist, ReducePartial, ReduceSplit, @@ -1362,6 +1363,126 @@ def test_reduce_unfactor_reduces(Sum, Product, backend: Backend): backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceUnfactor()) +def test_ground_cartesian_product_substitutes_dependent_plate_domain( + backend: Backend, +): + """Each grounded row value retains the domain of its plate assignment.""" + row_value, plate, t = backend.define_vars("row_value", "plate", "t", ret="scalar") + domain = backend.define_vars( + "domain", arg_types=(backend.scalar_typ,), ret="stream" + ) + factor = backend.define_vars( + "factor", + arg_types=( + backend.scalar_typ, + backend.scalar_typ, + backend.scalar_typ, + ), + ret="scalar", + ) + row = Operation.define(Mapping[tuple, backend.scalar_typ], name="row") # type: ignore[name-defined] + + lhs = Sum.reduce( + Product.reduce( + factor(t(), row()[(t(),)], row()[(t() + 1,)]), + {t: range(2)}, + ), + { + row: CartesianProduct.reduce( + Union.reduce( + [as_dict(((plate(),), row_value()))], + {row_value: domain(plate())}, + ), + {plate: range(3)}, + ) + }, + ) + + grounded = [Operation.define(row_value, name=f"row_value_{i}") for i in range(3)] + rhs = Sum.reduce( + Product.plus( + factor(0, grounded[0](), grounded[1]()), + factor(1, grounded[1](), grounded[2]()), + ), + {grounded[i]: domain(i) for i in range(3)}, + ) + + backend.check_rewrite( + lhs=lhs, + rhs=rhs, + rule=ReduceGroundCartesianProduct(), + ) + + +def test_ground_cartesian_product_declines_diagonal_chain(backend: Backend): + """A diagonal does not consume the off-diagonal cells of a two-dimensional row.""" + row_value, p, q, t = backend.define_vars("row_value", "p", "q", "t", ret="scalar") + factor = backend.define_vars( + "factor", + arg_types=(backend.scalar_typ, backend.scalar_typ), + ret="scalar", + ) + row = Operation.define(Mapping[tuple, backend.scalar_typ], name="row") # type: ignore[name-defined] + + lhs = Sum.reduce( + Product.reduce( + factor( + row()[(t(), t())], + row()[(t() + 1, t() + 1)], + ), + {t: range(1)}, + ), + { + row: CartesianProduct.reduce( + Union.reduce( + [as_dict(((p(), q()), row_value()))], {row_value: range(2)} + ), + {p: range(2), q: range(2)}, + ) + }, + ) + + backend.check_rewrite( + lhs=lhs, + rhs=lhs, + rule=ReduceGroundCartesianProduct(), + ) + + +def test_ground_cartesian_product_extra_value_streams(backend: Backend): + """Extra row value streams are retained.""" + row_value, row_value1, row_value2, p, q, t = backend.define_vars( + "row_value", "row_value1", "row_value2", "p", "q", "t", ret="scalar" + ) + row = Operation.define(Mapping[tuple, backend.scalar_typ], name="row") # type: ignore[name-defined] + + lhs = Sum.reduce( + Product.reduce(row()[(t(),)], {t: range(2)}), + { + row: CartesianProduct.reduce( + Union.reduce( + [as_dict(((p(),), row_value()))], {row_value: range(2), q: range(3)} + ), + {p: range(2)}, + ) + }, + ) + + rhs = Sum.reduce( + Product.plus(row_value1(), row_value2()), + { + row_value1: Union.reduce([row_value()], {q: range(3), row_value: range(2)}), + row_value2: Union.reduce([row_value()], {q: range(3), row_value: range(2)}), + }, + ) + + backend.check_rewrite( + lhs=lhs, + rhs=rhs, + rule=ReduceGroundCartesianProduct(), + ) + + @pytest.mark.parametrize("T,K", [(10, 3)]) def test_ground_cartesian_product_chain(T, K): """A chain over a cartesian product grounds into a forward pass.