Add generator-comprehension syntax for Monoid

[eb8680]

Gives Monoid a comprehension syntax, so a loop nest can be written the way you'd write it in Python and still be reduced symbolically.

Sum(f(x) * g(x, y) for x in xs for y in ys(x))

desugars to

Sum.reduce(f(x()) * g(x(), y()), {x: xs, y: ys(x())})

where x and y are fresh Operations standing for "an element of xs" and "an element of ys(x())".

How it works

Monoid.__call__ hands the generator object to effectful/internals/comprehension.py, which recovers the comprehension's syntax with the disassembler from the base PR and replays it against the monoid.

The interesting part is the typing, not the rewriting. Each loop target becomes an Operation whose return type is the element type of its stream, and that type has to be known before the target can be applied to anything in the body. Streams may also depend on earlier targets, as ys(x()) does, so element types are inferred one generator at a time, left to right: evaluate a stream, infer its element type, mint the target operation, bind it, then move on.

Some of what a comprehension may contain is syntax rather than operations. and, or, not and conditional expressions all ask their operands for a concrete bool, which a Term cannot supply, so they are rewritten into ite, which yields one of its arms outright when the condition is concrete and a term when it is not. Comparisons are left alone: numeric terms already implement == and < symbolically.

Changes outside the two new files

Three small, independent hunks, each reviewable on its own:

1. Monoid.__call__ (effectful/ops/monoid.py) — the entry point, with a doctest. This is the only part of the diff that is the feature itself.

2. ReduceDisequalityMask (effectful/ops/monoid.py) — an independent bugfix. It used or to pick the first non-None of two _neq_to_plus results, which asks a Term for its truthiness. Replaced with an explicit is None check, which is what the following if ret is not None already expected.

3. _jax_args (effectful/handlers/jax/monoid.py) — moved to Don't route all-scalar monoid ops through jax #729, which
   this branch is now stacked behind. These tests depend on it: jax.typing.ArrayLike
   is a union that includes bool/int/float/complex, and the JAX handlers extend
   the global EvaluateIntp at import time, so once anything imported
   effectful.handlers.jax.monoid a scalar comprehension body was silently narrowed to
   a float32 array. Six tests here (including two doctests) failed in a full-suite run
   but passed in isolation, purely on import order.

Nothing else from effectful/ops/monoid.py in #724 is included here — no ReduceGroundCartesianProduct, no inversion rewrite, no ReducePartial.unrolled. Those are separate branches.

Tests

tests/test_internals_comprehension.py: 734 passed, 5 xfailed.

The 5 xfails come from #724 as written and are not waiting on any sibling PR: a generator expression used as an inner stream (for y in (z for z in range(x))) raises NotImplementedError, marked strict=True. A list comprehension in the same position (for y in [z for z in range(2)]) does work and is tested.

Full suite (effectful/ tests/, excluding the LLM handler tests): 19704 passed, 2 skipped, 2085 xfailed, no failures. Doctests are on repo-wide via --doctest-modules, and the new Monoid.__call__ doctest passes both standalone and in a full run. ruff check and ruff format --diff are clean. mypy reports one pre-existing error in effectful/handlers/jax/monoid.py, present on staging-weighted and unrelated.

Split out of #724 for review. Stacked on #725, which is stacked on #729.
Review #729 → #725 → #727 in that order.

🤖 Generated with Claude Code