Simplify cartesian-product inversion with a peeled row - #730
Open
eb8680 wants to merge 15 commits into
Open
Conversation
jfeser
requested changes
Jul 28, 2026
| PlusEmpty(), | ||
| PlusSingle(), | ||
| PlusAssoc(), | ||
| ReduceUnfactor(), |
Contributor
There was a problem hiding this comment.
Can ReduceUnfactor be removed in this PR?
Comment on lines
+846
to
+850
| def __iter__(self): | ||
| return iter(()) | ||
|
|
||
| def __len__(self) -> int: | ||
| return 0 |
Contributor
There was a problem hiding this comment.
These should raise, I think. I assume they're unused.
| # what ``ReduceUnfactor`` does, applied here to the body's own factors | ||
| # rather than by renormalizing the whole candidate: the shape being | ||
| # looked for is a product of factors, so only the top level matters. | ||
| if isinstance(body, Term) and _is_monoid_plus(body.op): |
Contributor
There was a problem hiding this comment.
This shouldn't apply to any Monoid.plus.
| # Absorbing their streams back into this one exposes it again. This is | ||
| # what ``ReduceUnfactor`` does, applied here to the body's own factors | ||
| # rather than by renormalizing the whole candidate: the shape being | ||
| # looked for is a product of factors, so only the top level matters. |
Contributor
There was a problem hiding this comment.
I'm not convinced that doing one step of unfactoring is sufficient. What prevents a stream from being moved more deeply into the term?
Contributor
Author
There was a problem hiding this comment.
I can just reinstate the previous unfactoring behavior, I think this part is orthogonal to the other changes in the PR.
* 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
* 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
* 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
* 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
* 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
* 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
* 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
* 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
* 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
`_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) <noreply@anthropic.com>
* 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) <noreply@anthropic.com> * 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) <noreply@anthropic.com> * 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) <noreply@anthropic.com>
Standing a `_PeelRow` in for the row variable puts it in the mapping position of every `row[..]`, so a subscript is recognised by what it is indexing rather than by comparing against the variable it came from. One evaluation pass peels every summand and records the plate each folds over, replacing the per-summand `row_substitute` closure and its `InvalidIndexError` control flow. Nested same-monoid reductions are absorbed inline at the top level rather than by renormalizing the whole candidate under a dedicated interpretation, so `CartesianProductNormalizeIntp` is no longer needed. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
jfeser
force-pushed
the
eb-peel-row-inversion
branch
from
July 30, 2026 17:59
c62a3cb to
a0306a8
Compare
jfeser
force-pushed
the
staging-weighted
branch
from
August 10, 2026 16:49
42a4771 to
19e211f
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Rewrites the cartesian-product inversion rule
ReduceDistributeCartesianProductineffectful/ops/monoid.py. No behaviour change is intended.What changed
Peeling replaces the per-summand substitution closure.
_PeelRowis aMappingthat stands in for the row variable during oneevaluatepass. Because it occupies the mapping position of everyrow[..], a subscript is recognised by what it is indexing rather than by comparing the mapping against the variable it came from._peel_subscriptroutes_MappingTerm.__getitem__back to it (row indices are symbolic, so_MappingTerm.__getitem__would never hand a mapping a symbolic key on its own). Each subscript returns the narrowed row and records the plate variable used at the peeled position;row[p]with nothing left over returns the row's own value.This replaces
row_substituteplus itsInvalidIndexErrorcontrol flow: one pass peels every summand, and the per-summand checks that remain are just "the plate this summand's bundle binds is one of the plates the row was indexed by, and its range matches"._CannotPeelcarries the give-up signal, asInvalidIndexErrordid._PeelRowalso counts uses against peels per plate, sosurvives()can tell whether a plate is still referred to after peeling; only surviving plates need renaming onto the shared plate variable. The peeled row is bound to a freshpeeled_var = defop(stream_key)rather than reusingstream_key, since the row it names is a narrower one.The trailing
_to_bodyspecial case is gone: in the fully-peeled case the peeled row has already answered every subscript with the row's value, somonoid.reduce(combined, union_streams)is the same expression the oldhandler({_MappingTerm.__getitem__: _to_body})(evaluate)(combined)produced.CartesianProductNormalizeIntpis removed.Factormay have moved product factors into nested reductions of the same monoid, hiding the product this rule looks for. Previously the whole candidate was rebuilt under a dedicatedCartesianProductNormalizeIntpso thatReduceUnfactorcould merge the stream bundles. That is now done inline over the body's own top-level factors: the shape being looked for is a product of factors, so only the top level matters. The nested streams are absorbed intostreamswhen they do not collide, which is exactly whatReduceUnfactorwould have done.grep -rn CartesianProductNormalizeIntpover the repo (excluding.git) matched only its definition and its single use insidemonoid.py— nothing ineffectful/handlers/jax/monoid.py,tests/, ordocs/referred to it. Its removal also drops the now-unusedfrom effectful.internals.runtime import interpreterimport;defopis now imported fromeffectful.ops.syntax.Testing
No new tests:
tests/test_ops_monoid.pyandtests/test_handlers_jax_monoid.pyalready cover this rule (530 tests, all passing). A fulltests/sweep excludingtests/test_handlers_llm_*.pygives 18317 passed, 2 skipped, 2078 xfailed, 0 failed, matchingstaging-weighted.mypy,ruff checkandruff formatare clean apart from the pre-existingeffectful/handlers/jax/monoid.py:382error, which is also present onstaging-weighted.Note for review
Six
typing.cast(...)calls in the new code are not in #724 — they are runtime no-ops added purely to keepmypyclean, sinceTerm.argselements type asExpr. They are the only difference between this rule's body here and the one in #724, which is otherwise verbatim.Split out of #724 for review. Independent of the sibling
ReduceGroundCartesianProductPR, also split out of #724.🤖 Generated with Claude Code