Skip to content

Add ReduceGroundCartesianProduct - #731

Open
eb8680 wants to merge 16 commits into
staging-weightedfrom
eb-ground-cartesian-product
Open

Add ReduceGroundCartesianProduct#731
eb8680 wants to merge 16 commits into
staging-weightedfrom
eb-ground-cartesian-product

Conversation

@eb8680

@eb8680 eb8680 commented Jul 28, 2026

Copy link
Copy Markdown
Contributor

Adds ReduceGroundCartesianProduct, the fallback rule for a reduction over a
cartesian-product row stream whose body admits no per-plate factorization.

The rewrite

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.

_GroundRow is a Mapping that stands in for the row variable. Unrolling the
body's plate folds hands it concrete keys, and it answers each with a variable
minted on first sight (the same variable for a repeated key, which is what ties
the factors of a chain together). A key that is still symbolic never reaches it:
the subscript stays a term until the fold supplying the index is expanded.
_GroundRow.residual returns whatever plate values no key consumed, which stay
in the expression as a smaller cartesian product; sealed + fvsof reject a
half-ground result, in which case the rule declines.

When this applies vs. ReduceDistributeCartesianProduct

Where the body folds over the plates uniformly, ReduceDistributeCartesianProduct
inverts the reduction instead — cheaper, and it leaves the plate fold intact.
Grounding is for bodies where inversion does not apply, e.g. a chain
X[t], X[t+1] couples adjacent plate indices. Grounding costs one variable per
assignment (|P| variables over D) instead of enumerating the |D|^|P| rows,
after which what is left is an ordinary variable-elimination problem that
Factor solves. ReduceGroundCartesianProduct() is registered in
NormalizeIntp immediately before ReduceDistributeCartesianProduct().

ReduceDistributeCartesianProduct itself is untouched here.

ReducePartial.unrolled

ReducePartial gains an optional unrolled set of stream variables and skips
any stream not in it. Grounding installs ReducePartial(plate_vars) so that
expanding the body expands only the plate folds it is grounding over, and not,
say, the value domains.

Test

tests/test_ops_monoid.py::test_ground_cartesian_product_chain reduces a chain
Sum_ixs Prod_t phi[t][ixs[t]][ixs[t+1]] and checks the result against the
forward algorithm computed in plain Python.

It is written longhand with explicit defop variables. In #724 it is written
as

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))
)

using the generator-comprehension syntax that lands in a later PR; the body
below is exactly what desugar_comprehension produces for it.

T is 10 rather than the 20 used in #724: normalization is superlinear in the
chain length until the separate evaluate performance fix lands, and this
branch is deliberately not stacked on it. At T = 10 the test takes ~6s; at
T = 20 it takes over an hour. K = 3 as in #724.

The test does discriminate: unregistering the rule makes normalization a no-op
and the added structural assertion fails immediately. (Without that assertion
the test would still pass at T = 10, just 11x slower, because evaluation
falls back to enumerating all K ** T rows — which is why the assertion is
there.)

Prerequisite: #729

This PR is stacked on #729 (eb-jax-scalar-plus), which fixes _jax_args in
effectful/handlers/jax/monoid.py. That fix is a genuine prerequisite here, not
a cosmetic one: jax.typing.ArrayLike includes bool/int/float/complex,
so the jax handlers claimed pure-Python scalar arithmetic. Without it the new
test does not merely produce a different type — it raises TypeError: Cannot unify type <class 'int'> with <class 'jax.Array'>
, because
tests/test_ops_monoid.py imports effectful.handlers.jax.monoid
unconditionally.

This branch's own diff is just effectful/ops/monoid.py and
tests/test_ops_monoid.py.

The ReduceDisequalityMask or -> is None fix that also lives in #724 was
checked and is not needed here: the full suite passes without it.

Testing

  • tests/test_ops_monoid.py + tests/test_handlers_jax_monoid.py: 531 passed.
  • Full tests/ sweep excluding tests/test_handlers_llm_*.py: 18318 passed,
    2 skipped, 2078 xfailed — versus 18317/2/2078 on pristine staging-weighted,
    the difference being the new test. No regressions.
  • ruff check, ruff format --diff clean. mypy clean apart from the one
    pre-existing effectful/handlers/jax/monoid.py error already on
    staging-weighted.

Split out of #724 for review. Stacked on #729. Independent of the sibling
ReduceDistributeCartesianProduct refactor (#730), also split out of #724.

🤖 Generated with Claude Code

@eb8680
eb8680 force-pushed the eb-ground-cartesian-product branch from c5efdf2 to 93a5895 Compare July 28, 2026 16:07
@eb8680
eb8680 changed the base branch from staging-weighted to eb-jax-scalar-plus July 28, 2026 16:07
@eb8680 eb8680 added the blocked label Jul 28, 2026
Base automatically changed from eb-jax-scalar-plus to staging-weighted July 28, 2026 16:19
@jfeser jfeser removed the blocked label Jul 28, 2026
@eb8680
eb8680 requested a review from jfeser July 28, 2026 22:23
@eb8680
eb8680 marked this pull request as ready for review July 29, 2026 00:08
@eb8680

eb8680 commented Jul 29, 2026

Copy link
Copy Markdown
Contributor Author

The slop in the PR description is from disaggregating #724. The code in this PR was largely written and checked by me.

@jfeser jfeser left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I pushed some new tests that currently fail. We should either fix those cases or forward. The diagonal stream consumption issue also affects ReduceDistributeCartesianProduct (there it causes a performance problem, not a correctness one).

jfeser and others added 14 commits July 30, 2026 11:43
* 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>
@jfeser
jfeser force-pushed the staging-weighted branch from 0553c4a to fe1e760 Compare July 30, 2026 17:05
eb8680 and others added 2 commits July 30, 2026 14:58
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) <noreply@anthropic.com>
@jfeser
jfeser force-pushed the eb-ground-cartesian-product branch from 7acbe8c to 7be9b05 Compare July 30, 2026 18:58
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants