From 523ff77b4d15d50bf8e3d70af1e5ea4e4060f4ec Mon Sep 17 00:00:00 2001
From: Jack Feser <jack.feser@gmail.com>
Date: Wed, 29 Jul 2026 09:31:59 -0400
Subject: [PATCH 01/17] wip

---
 effectful/ops/monoid.py  | 61 +++++++++++++++++++++++++++++++++++++++-
 tests/test_ops_monoid.py | 53 ++++++++++++++++++++++++++++++++++
 2 files changed, 113 insertions(+), 1 deletion(-)

diff --git a/effectful/ops/monoid.py b/effectful/ops/monoid.py
index 37ae0bf6..8f4fd9b6 100644
--- a/effectful/ops/monoid.py
+++ b/effectful/ops/monoid.py
@@ -203,6 +203,31 @@ def __init__(self, name: str, identity: T, zero: T):
         self.zero = zero
 
 
+class ProductMonoid[L, R](Monoid[tuple[L, R]]):
+    """The componentwise product of two monoids.
+
+    ``ProductMonoid(left, right).plus`` combines the first components with
+    ``left.plus`` and the second components with ``right.plus``.
+    """
+
+    left: Monoid[L]
+    right: Monoid[R]
+
+    def __init__(self, left: Monoid[L], right: Monoid[R], name: str | None = None):
+        self.left = left
+        self.right = right
+        super().__init__(
+            name=name or f"{left.__name__}×{right.__name__}",
+            identity=(left.identity, right.identity),
+        )
+
+        # Product operations inherit these algebraic properties componentwise.
+        if is_commutative(left) and is_commutative(right):
+            is_commutative.register(self)
+        if is_idempotent(left) and is_idempotent(right):
+            is_idempotent.register(self)
+
+
 Min = Monoid(name="Min", identity=float("inf"))
 Max = Monoid(name="Max", identity=-float("inf"))
 ArgMin = Monoid(name="ArgMin", identity=(Min.identity, None))
@@ -1613,6 +1638,39 @@ def _disjoint_merge[K, V](*dicts: Mapping[K, V]) -> Mapping[K, V]:
     return merged
 
 
+class AssignmentPlus(ObjectInterpretation):
+    """Disjoint-union implementation of :data:`Assignment`."""
+
+    @implements(Assignment.plus)
+    def plus(self, *args):
+        if not args:
+            return Assignment.identity
+        if any(isinstance(arg, Term) for arg in args):
+            return fwd()
+        if not all(isinstance(arg, Mapping) for arg in args):
+            return fwd()
+        return _disjoint_merge(*args)
+
+
+class ProductMonoidPlus(ObjectInterpretation):
+    """Componentwise implementation of :class:`ProductMonoid`."""
+
+    @implements(Monoid.plus)
+    def plus(self, monoid, *args):
+        if not isinstance(monoid, ProductMonoid):
+            return fwd()
+        if not args:
+            return monoid.identity
+        if any(isinstance(arg, Term) for arg in args):
+            return fwd()
+        if not all(isinstance(arg, tuple) and len(arg) == 2 for arg in args):
+            return fwd()
+        return (
+            monoid.left.plus(*(arg[0] for arg in args)),
+            monoid.right.plus(*(arg[1] for arg in args)),
+        )
+
+
 class CartesianProductPlus(ObjectInterpretation):
     """Pure-Python implementation of :data:`CartesianProduct`."""
 
@@ -2130,10 +2188,11 @@ def extend(self, *intps: Interpretation) -> typing.Self:
     ReducePartial(),
     DeltaConcrete(),
     SumPlus(),
-    SumInverse(),
     MinPlus(),
     MaxPlus(),
     ProductPlus(),
+    AssignmentPlus(),
+    ProductMonoidPlus(),
     ArgMinPlus(),
     ArgMaxPlus(),
     CartesianProductPlus(),
diff --git a/tests/test_ops_monoid.py b/tests/test_ops_monoid.py
index 6bc4b881..6d28f002 100644
--- a/tests/test_ops_monoid.py
+++ b/tests/test_ops_monoid.py
@@ -1,5 +1,6 @@
 import functools
 import math
+import numbers
 import operator
 import sys
 import typing
@@ -12,6 +13,7 @@
 import effectful.handlers.jax.monoid  # noqa: F401
 from effectful.ops.monoid import (
     And,
+    Assignment,
     CartesianProduct,
     CartesianProductPlus,
     EliminateSingletonStreams,
@@ -39,6 +41,7 @@
     PlusPartial,
     PlusSingle,
     Product,
+    ProductMonoid,
     ReduceDependentRangeMask,
     ReduceDisjunctiveDisequalityMask,
     ReduceDistributeCartesianProduct,
@@ -568,6 +571,36 @@ def test_plus_zero(monoid, backend: Backend):
     backend.check_rewrite(lhs=lhs_left, rhs=rhs, rule={})
 
 
+def test_assignment_plus_disjoint_merge():
+    with handler(EvaluateIntp):
+        assert Assignment.plus({"x": 1}, {"y": 2}) == {"x": 1, "y": 2}
+        assert Assignment.plus() == {}
+
+
+def test_assignment_plus_rejects_duplicate_keys():
+    with handler(EvaluateIntp), pytest.raises(ValueError, match="Duplicate key"):
+        Assignment.plus({"x": 1}, {"x": 1})
+
+
+def test_product_monoid_plus():
+    scored_assignment = ProductMonoid(Sum, Assignment)
+
+    with handler(EvaluateIntp):
+        result = scored_assignment.plus((2, {"x": 1}), (3, {"y": 2}))
+
+    assert result == (5, {"x": 1, "y": 2})
+    assert scored_assignment.identity == (0, {})
+    assert is_commutative(scored_assignment)
+    assert not is_idempotent(scored_assignment)
+
+
+def test_product_monoid_inherits_idempotence():
+    extrema = ProductMonoid(Min, Max)
+
+    assert is_commutative(extrema)
+    assert is_idempotent(extrema)
+
+
 def test_plus_partial():
     backend = IntBackend()
     x = backend.define_vars("x", ret="scalar")
@@ -1601,3 +1634,23 @@ 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())
+
+
+def test_reduce_argmin(backend: Backend):
+    x, y, z = backend.define_vars("x", "y", "z", ret="scalar")
+    X, Y, Z = backend.define_vars("X", "Y", "Z", ret="stream")
+
+    class ArgValue[T: numbers.Number]:
+        score: T = Sum.identity
+        assignment: Mapping[Operation, Any]
+
+    ArgSum = Monoid(name="ArgSum", identity=ArgValue())
+
+    lhs = Min.reduce(
+        (x() - 1) ** 2,
+        {
+            x: ArgSum.weighted(
+                range(3),
+            )
+        },
+    )

From e99a96088f6a07b4785a1189d11056be8e109455 Mon Sep 17 00:00:00 2001
From: Jack Feser <jack.feser@gmail.com>
Date: Wed, 29 Jul 2026 12:03:01 -0400
Subject: [PATCH 02/17] wip

---
 effectful/ops/monoid.py  | 88 ++++++++++++++++++++++++++++++++++++++++
 tests/test_ops_monoid.py | 30 +++++++-------
 2 files changed, 104 insertions(+), 14 deletions(-)

diff --git a/effectful/ops/monoid.py b/effectful/ops/monoid.py
index 8f4fd9b6..d5244fbe 100644
--- a/effectful/ops/monoid.py
+++ b/effectful/ops/monoid.py
@@ -203,6 +203,42 @@ def __init__(self, name: str, identity: T, zero: T):
         self.zero = zero
 
 
+class Optimum[T]:
+    """A value together with an assignment that attains it."""
+
+    __slots__ = ("assignment", "value")
+
+    def __init__(self, value: T, assignment: Mapping[Operation, Any] | None):
+        self.value = value
+        self.assignment = assignment
+
+    @property
+    def feasible(self) -> bool:
+        return self.assignment is not None
+
+    def __eq__(self, other: object) -> bool:
+        return (
+            isinstance(other, Optimum)
+            and self.value == other.value
+            and self.assignment == other.assignment
+        )
+
+    def __repr__(self) -> str:
+        return f"Optimum(value={self.value!r}, assignment={self.assignment!r})"
+
+
+@evaluate.register(Optimum)
+def _evaluate_optimum(expr: Optimum, **kwargs) -> Optimum:
+    """Evaluate an optimum's score and assignment values, preserving variable keys."""
+
+    assignment = (
+        None
+        if expr.assignment is None
+        else {variable: evaluate(value) for variable, value in expr.assignment.items()}
+    )
+    return Optimum(evaluate(expr.value), assignment)
+
+
 class ProductMonoid[L, R](Monoid[tuple[L, R]]):
     """The componentwise product of two monoids.
 
@@ -1612,6 +1648,24 @@ def plus(self, *args):
         return min(args, key=lambda a: a[0])
 
 
+class MinOptimumPlus(ObjectInterpretation):
+    """Lift :data:`Min` to values carrying minimizing assignments."""
+
+    @implements(Min.plus)
+    def plus(self, *args):
+        if not args or not all(isinstance(a, Optimum) for a in args):
+            return fwd()
+        feasible = [a for a in args if a.feasible]
+        if not feasible:
+            return Optimum(Min.identity, None)
+        if any(isinstance(a.value, Term) for a in feasible):
+            return fwd()
+        try:
+            return min(feasible, key=lambda a: a.value)
+        except (TypeError, ValueError):
+            return fwd()
+
+
 class ArgMaxPlus(ObjectInterpretation):
     """Scalar score implementation of :data:`ArgMax`."""
 
@@ -1777,6 +1831,35 @@ def plus(self, monoid, *args):
         return fwd()
 
 
+class SumOptimumPlus(ObjectInterpretation):
+    """Lift :data:`Sum` to values carrying minimizing assignments."""
+
+    @implements(Sum.plus)
+    def plus(self, *args):
+        if not args or not all(isinstance(arg, Optimum) for arg in args):
+            return fwd()
+        if any(not arg.feasible for arg in args):
+            return Optimum(float("inf"), None)
+        return Optimum(
+            Sum.plus(*(arg.value for arg in args)),
+            _disjoint_merge(*(typing.cast(Mapping, arg.assignment) for arg in args)),
+        )
+
+
+class PlusCastOptimum(ObjectInterpretation):
+    """Promote plain scores when adding an assignment-carrying weight."""
+
+    @implements(Sum.plus)
+    def plus(self, *args):
+        if any(isinstance(arg, Optimum) for arg in args) and any(
+            not isinstance(arg, Optimum) for arg in args
+        ):
+            return Sum.plus(
+                *(arg if isinstance(arg, Optimum) else Optimum(arg, {}) for arg in args)
+            )
+        return fwd()
+
+
 class EliminateSingletonStreams(ObjectInterpretation):
     """Eliminate a length-1 stream by substituting its sole element.
 
@@ -2188,7 +2271,10 @@ def extend(self, *intps: Interpretation) -> typing.Self:
     ReducePartial(),
     DeltaConcrete(),
     SumPlus(),
+    SumOptimumPlus(),
+    PlusCastOptimum(),
     MinPlus(),
+    MinOptimumPlus(),
     MaxPlus(),
     ProductPlus(),
     AssignmentPlus(),
@@ -2249,6 +2335,8 @@ def extend(self, *intps: Interpretation) -> typing.Self:
     PlusOrder(),
     PlusCastFloat(),
     PlusCastIterable(),
+    SumOptimumPlus(),
+    PlusCastOptimum(),
     MaskFusion(),
     MaskBool(),
     WhereHoist(),
diff --git a/tests/test_ops_monoid.py b/tests/test_ops_monoid.py
index 6d28f002..f11264bd 100644
--- a/tests/test_ops_monoid.py
+++ b/tests/test_ops_monoid.py
@@ -1,6 +1,5 @@
 import functools
 import math
-import numbers
 import operator
 import sys
 import typing
@@ -30,6 +29,7 @@
     MonoidOverMapping,
     MonoidOverSequence,
     NormalizeIntp,
+    Optimum,
     Or,
     PlusAssoc,
     PlusCastIterable,
@@ -1637,20 +1637,22 @@ def test_reduce_unfactor_reduces(Sum, Product, backend: Backend):
 
 
 def test_reduce_argmin(backend: Backend):
-    x, y, z = backend.define_vars("x", "y", "z", ret="scalar")
-    X, Y, Z = backend.define_vars("X", "Y", "Z", ret="stream")
-
-    class ArgValue[T: numbers.Number]:
-        score: T = Sum.identity
-        assignment: Mapping[Operation, Any]
+    x = backend.define_vars("x", ret="scalar")
 
-    ArgSum = Monoid(name="ArgSum", identity=ArgValue())
+    def record_assignment(value):
+        return Optimum(Sum.identity, {x: value})
 
-    lhs = Min.reduce(
+    expr = Min.reduce(
         (x() - 1) ** 2,
-        {
-            x: ArgSum.weighted(
-                range(3),
-            )
-        },
+        {x: Sum.weighted(range(3), record_assignment)},
     )
+
+    with handler(NormalizeIntp):
+        norm_expr = evaluate(expr)
+
+    breakpoint()
+
+    with handler(NormalizeIntp), handler(EvaluateIntp):
+        result = evaluate(expr)
+
+    assert result == Optimum(0, {x: 1})

From 3b2e4b3adbf8725bf7a6355a95a19403b8ede0a7 Mon Sep 17 00:00:00 2001
From: Jack Feser <jack.feser@gmail.com>
Date: Wed, 29 Jul 2026 15:24:37 -0400
Subject: [PATCH 03/17] add argmin tests

---
 tests/test_ops_monoid.py | 43 ++++++++++++++++++++++++++++++----------
 1 file changed, 33 insertions(+), 10 deletions(-)

diff --git a/tests/test_ops_monoid.py b/tests/test_ops_monoid.py
index f11264bd..b4f9f887 100644
--- a/tests/test_ops_monoid.py
+++ b/tests/test_ops_monoid.py
@@ -68,7 +68,7 @@
     solve_group_equality,
 )
 from effectful.ops.semantics import coproduct, evaluate, fvsof, handler
-from effectful.ops.syntax import as_dict, ite, range_, syntactic_eq
+from effectful.ops.syntax import as_dict, deffn, ite, range_, syntactic_eq
 from effectful.ops.types import NotHandled, Operation, Term
 from tests._monoid_helpers import Backend, IntBackend, JaxBackend, syntactic_eq_alpha
 
@@ -1637,22 +1637,45 @@ def test_reduce_unfactor_reduces(Sum, Product, backend: Backend):
 
 
 def test_reduce_argmin(backend: Backend):
-    x = backend.define_vars("x", ret="scalar")
-
-    def record_assignment(value):
-        return Optimum(Sum.identity, {x: value})
+    x, v = backend.define_vars("x", "v", ret="scalar")
 
     expr = Min.reduce(
         (x() - 1) ** 2,
-        {x: Sum.weighted(range(3), record_assignment)},
+        {x: Sum.weighted(range(3), deffn(Optimum(Sum.identity, {x: v()}), v))},
     )
 
-    with handler(NormalizeIntp):
-        norm_expr = evaluate(expr)
+    with handler(NormalizeIntp), handler(EvaluateIntp):
+        result = evaluate(expr)
+
+    assert result == Optimum(0, {x: 1})
+
+
+def test_reduce_argmin_sum(backend: Backend):
+    x, v = backend.define_vars("x", "v", ret="scalar")
 
-    breakpoint()
+    expr = Min.reduce(
+        Sum.plus((x() - 1) ** 2, 2 * (x() - 2) ** 2),
+        {x: Sum.weighted(range(3), deffn(Optimum(Sum.identity, {x: v()}), v))},
+    )
 
     with handler(NormalizeIntp), handler(EvaluateIntp):
         result = evaluate(expr)
 
-    assert result == Optimum(0, {x: 1})
+    assert result == Optimum(1, {x: 2})
+
+
+def test_reduce_argmin_sum_disjoint(backend: Backend):
+    x, y, v = backend.define_vars("x", "y", "v", ret="scalar")
+
+    expr = Min.reduce(
+        Sum.plus((x() - 1) ** 2, (y() - 2) ** 2),
+        {
+            x: Sum.weighted(range(3), deffn(Optimum(Sum.identity, {x: v()}), v)),
+            y: Sum.weighted(range(3), deffn(Optimum(Sum.identity, {y: v()}), v)),
+        },
+    )
+
+    with handler(NormalizeIntp), handler(EvaluateIntp):
+        result = evaluate(expr)
+
+    assert result == Optimum(0, {x: 1, y: 2})

From 1457a4d9f4ee8e4f0c0f9722085be5436dedda3b Mon Sep 17 00:00:00 2001
From: Jack Feser <jack.feser@gmail.com>
Date: Mon, 10 Aug 2026 18:55:54 -0400
Subject: [PATCH 04/17] add jax versions

---
 effectful/handlers/jax/monoid.py  | 102 ++++++++++++++++++++++++++++++
 effectful/ops/monoid.py           |  19 ++++++
 tests/test_handlers_jax_monoid.py |  45 +++++++++++++
 3 files changed, 166 insertions(+)

diff --git a/effectful/handlers/jax/monoid.py b/effectful/handlers/jax/monoid.py
index 28233e76..6f353e92 100644
--- a/effectful/handlers/jax/monoid.py
+++ b/effectful/handlers/jax/monoid.py
@@ -25,6 +25,7 @@
     Min,
     Monoid,
     NormalizeIntp,
+    Optimum,
     Or,
     Product,
     Streams,
@@ -51,6 +52,25 @@
 logger = logging.getLogger(__name__)
 
 
+def _optimum_flatten(optimum: Optimum):
+    keys = None if optimum.assignment is None else tuple(optimum.assignment)
+    values = () if optimum.assignment is None else tuple(optimum.assignment.values())
+    return (optimum.value, *values), keys
+
+
+def _optimum_unflatten(keys, children):
+    value, *assignment_values = children
+    assignment = (
+        None if keys is None else dict(zip(keys, assignment_values, strict=True))
+    )
+    return Optimum(value, assignment)
+
+
+# ``Optimum`` lives in the backend-independent module, but once the JAX backend
+# is imported it should be a valid input/output of ``jax.jit``.
+jax.tree_util.register_pytree_node(Optimum, _optimum_flatten, _optimum_unflatten)
+
+
 is_equality.register(jnp.equal)
 for a, b in {
     (jnp.less, jnp.greater),
@@ -259,6 +279,87 @@ def __call__(
 ARRAY_REDUCTORS[LogSumExp] = logsumexp
 
 
+class ReduceOptimum(ObjectInterpretation):
+    """Reduce an assignment-carrying JAX score with ``argmin`` or ``argmax``.
+
+    This is deliberately a lowering for the normalized ``Optimum`` body rather
+    than for weighted-stream syntax.  ``ReduceWeightedStream`` turns the latter
+    into this form, so the same kernel handles both spellings.
+
+    The initial kernel supports independent ``range`` streams and assignments
+    that record stream variables directly.  Those are the normal form produced
+    by ``Sum.weighted(..., lambda v: Optimum(0, {x: v}))``.
+    """
+
+    @implements(Monoid.reduce)
+    def reduce(self, monoid, body, streams):
+        if monoid not in (Min, Max) or not isinstance(body, Optimum):
+            return fwd()
+        if not issubclass(typeof(body.value), jax.Array):
+            return fwd()
+        if not body.feasible or not isinstance(body.assignment, Mapping):
+            return fwd()
+        if not streams or not all(
+            isinstance(stream, range) for stream in streams.values()
+        ):
+            return fwd()
+
+        # For now assignments must be direct references to reduction variables.
+        # Keeping this check narrow is preferable to silently returning a wrong
+        # provenance value for an arbitrary assignment expression.
+        assignment_vars = {}
+        for key, value in body.assignment.items():
+            if not (
+                isinstance(value, Term)
+                and not value.args
+                and not value.kwargs
+                and value.op in streams
+            ):
+                return fwd()
+            assignment_vars[key] = value.op
+
+        if any(len(stream) == 0 for stream in streams.values()):
+            return Optimum(monoid.identity, None)
+
+        score_fvs = fvsof(body.value)
+        used = tuple(k for k in streams if k in score_fvs)
+        used_streams = {k: streams[k] for k in used}
+
+        if used:
+            # Materialize one leading positional axis per used stream.  Existing
+            # delta lowering performs vectorized substitution and preserves any
+            # trailing batch dimensions of the score.
+            score = monoid.reduce(
+                monoid.delta(tuple(k() for k in used), body.value), used_streams
+            )
+            if not isinstance(score, jax.Array | jax.core.Tracer):
+                return fwd()
+
+            reduction_shape = tuple(len(typing.cast(range, streams[k])) for k in used)
+            if tuple(score.shape[: len(used)]) != reduction_shape:
+                return fwd()
+
+            flat_size = functools.reduce(lambda a, b: a * b, reduction_shape, 1)
+            flat_score = jnp.reshape(score, (flat_size, *score.shape[len(used) :]))
+            arg_reduce = jnp.argmin if monoid is Min else jnp.argmax
+            flat_index = arg_reduce(flat_score, axis=0)
+            value = jnp.take_along_axis(flat_score, flat_index[None], axis=0)[0]
+            coordinates = jnp.unravel_index(flat_index, reduction_shape)
+            positions = dict(zip(used, coordinates, strict=True))
+        else:
+            # An invariant score ties everywhere; monoid reduction chooses the
+            # first candidate, matching both Python's min/max and JAX argmin/max.
+            value = body.value
+            positions = {}
+
+        assignment = {}
+        for key, variable in assignment_vars.items():
+            stream = typing.cast(range, streams[variable])
+            position = positions.get(variable, 0)
+            assignment[key] = stream.start + stream.step * position
+        return Optimum(value, assignment)
+
+
 class ReduceArray(ObjectInterpretation):
     """Reduce an array body over range streams."""
 
@@ -822,6 +923,7 @@ def einsum(
     ReduceDeltaSimpleRange(),
     ReduceArrayScan(),
     PlusCastArray(),
+    ReduceOptimum(),
 )
 
 NormalizeIntp.extend(
diff --git a/effectful/ops/monoid.py b/effectful/ops/monoid.py
index d5244fbe..23791d29 100644
--- a/effectful/ops/monoid.py
+++ b/effectful/ops/monoid.py
@@ -1666,6 +1666,24 @@ def plus(self, *args):
             return fwd()
 
 
+class MaxOptimumPlus(ObjectInterpretation):
+    """Lift :data:`Max` to values carrying maximizing assignments."""
+
+    @implements(Max.plus)
+    def plus(self, *args):
+        if not args or not all(isinstance(a, Optimum) for a in args):
+            return fwd()
+        feasible = [a for a in args if a.feasible]
+        if not feasible:
+            return Optimum(Max.identity, None)
+        if any(isinstance(a.value, Term) for a in feasible):
+            return fwd()
+        try:
+            return max(feasible, key=lambda a: a.value)
+        except (TypeError, ValueError):
+            return fwd()
+
+
 class ArgMaxPlus(ObjectInterpretation):
     """Scalar score implementation of :data:`ArgMax`."""
 
@@ -2276,6 +2294,7 @@ def extend(self, *intps: Interpretation) -> typing.Self:
     MinPlus(),
     MinOptimumPlus(),
     MaxPlus(),
+    MaxOptimumPlus(),
     ProductPlus(),
     AssignmentPlus(),
     ProductMonoidPlus(),
diff --git a/tests/test_handlers_jax_monoid.py b/tests/test_handlers_jax_monoid.py
index a538664d..dabe4c27 100644
--- a/tests/test_handlers_jax_monoid.py
+++ b/tests/test_handlers_jax_monoid.py
@@ -21,11 +21,15 @@
 from effectful.ops.monoid import (
     EliminateSingletonStreams,
     EvaluateIntp,
+    Max,
+    Min,
     NormalizeIntp,
+    Optimum,
     Product,
     Sum,
 )
 from effectful.ops.semantics import coproduct, evaluate, handler
+from effectful.ops.syntax import deffn
 from tests._monoid_helpers import JaxBackend
 
 MONOIDS = [
@@ -124,6 +128,47 @@ def test_reduce_array_2(monoid, reductor, backend: JaxBackend):
     assert jnp.allclose(actual, expected)
 
 
+@pytest.mark.parametrize(
+    "monoid,expected_value,expected_assignment",
+    [(Min, 0, 1), (Max, 4, 3)],
+)
+def test_reduce_optimum_weighted_stream(
+    monoid, expected_value, expected_assignment, backend: JaxBackend
+):
+    x, v = backend.define_vars("x", "v", ret="scalar")
+    expr = monoid.reduce(
+        (x() - 1) ** 2,
+        {x: Sum.weighted(range(4), deffn(Optimum(0, {x: v()}), v))},
+    )
+
+    with handler(NormalizeIntp), handler(EvaluateIntp):
+        actual = evaluate(expr)
+
+    assert isinstance(actual, Optimum)
+    assert jnp.array_equal(actual.value, expected_value)
+    assert actual.assignment is not None
+    assert jnp.array_equal(actual.assignment[x], expected_assignment)
+
+
+def test_reduce_optimum_jit(backend: JaxBackend):
+    x = backend.define_vars("x", ret="scalar")
+
+    @jax.jit
+    def run(scores):
+        with handler(NormalizeIntp), handler(EvaluateIntp):
+            return Min.reduce(
+                Optimum(unbind_dims(scores, x), {x: x()}),
+                {x: range(scores.shape[0])},
+            )
+
+    actual = run(jnp.asarray([3.0, 1.0, 2.0]))
+
+    assert isinstance(actual, Optimum)
+    assert jnp.array_equal(actual.value, 1.0)
+    assert actual.assignment is not None
+    assert jnp.array_equal(actual.assignment[x], 1)
+
+
 SCALAR_PLUS = [
     pytest.param(Sum, 3.0, id="Sum"),
     pytest.param(Product, 2.0, id="Product"),

From 536cd79c4a88e9dce85f3042374a48e784ffaf1e Mon Sep 17 00:00:00 2001
From: Jack Feser <jack.feser@gmail.com>
Date: Mon, 10 Aug 2026 18:57:22 -0400
Subject: [PATCH 05/17] simplify

---
 effectful/handlers/jax/monoid.py |  34 +-----
 effectful/ops/monoid.py          | 184 ++++++++-----------------------
 tests/test_ops_monoid.py         |  40 ++-----
 3 files changed, 58 insertions(+), 200 deletions(-)

diff --git a/effectful/handlers/jax/monoid.py b/effectful/handlers/jax/monoid.py
index 6f353e92..20433e3d 100644
--- a/effectful/handlers/jax/monoid.py
+++ b/effectful/handlers/jax/monoid.py
@@ -18,6 +18,7 @@
 from effectful.handlers.jax.scipy.special import logsumexp
 from effectful.ops.monoid import (
     And,
+    Body,
     CartesianProduct,
     EvaluateIntp,
     LogSumExp,
@@ -52,23 +53,9 @@
 logger = logging.getLogger(__name__)
 
 
-def _optimum_flatten(optimum: Optimum):
-    keys = None if optimum.assignment is None else tuple(optimum.assignment)
-    values = () if optimum.assignment is None else tuple(optimum.assignment.values())
-    return (optimum.value, *values), keys
-
-
-def _optimum_unflatten(keys, children):
-    value, *assignment_values = children
-    assignment = (
-        None if keys is None else dict(zip(keys, assignment_values, strict=True))
-    )
-    return Optimum(value, assignment)
-
-
 # ``Optimum`` lives in the backend-independent module, but once the JAX backend
 # is imported it should be a valid input/output of ``jax.jit``.
-jax.tree_util.register_pytree_node(Optimum, _optimum_flatten, _optimum_unflatten)
+jax.tree_util.register_dataclass(Optimum)
 
 
 is_equality.register(jnp.equal)
@@ -282,23 +269,17 @@ def __call__(
 class ReduceOptimum(ObjectInterpretation):
     """Reduce an assignment-carrying JAX score with ``argmin`` or ``argmax``.
 
-    This is deliberately a lowering for the normalized ``Optimum`` body rather
-    than for weighted-stream syntax.  ``ReduceWeightedStream`` turns the latter
-    into this form, so the same kernel handles both spellings.
-
     The initial kernel supports independent ``range`` streams and assignments
     that record stream variables directly.  Those are the normal form produced
     by ``Sum.weighted(..., lambda v: Optimum(0, {x: v}))``.
     """
 
     @implements(Monoid.reduce)
-    def reduce(self, monoid, body, streams):
+    def reduce(self, monoid: Monoid, body: Body, streams: Streams):
         if monoid not in (Min, Max) or not isinstance(body, Optimum):
             return fwd()
         if not issubclass(typeof(body.value), jax.Array):
             return fwd()
-        if not body.feasible or not isinstance(body.assignment, Mapping):
-            return fwd()
         if not streams or not all(
             isinstance(stream, range) for stream in streams.values()
         ):
@@ -309,17 +290,12 @@ def reduce(self, monoid, body, streams):
         # provenance value for an arbitrary assignment expression.
         assignment_vars = {}
         for key, value in body.assignment.items():
-            if not (
-                isinstance(value, Term)
-                and not value.args
-                and not value.kwargs
-                and value.op in streams
-            ):
+            if not (isinstance(value, Term) and value.op in streams):
                 return fwd()
             assignment_vars[key] = value.op
 
         if any(len(stream) == 0 for stream in streams.values()):
-            return Optimum(monoid.identity, None)
+            return monoid.identity
 
         score_fvs = fvsof(body.value)
         used = tuple(k for k in streams if k in score_fvs)
diff --git a/effectful/ops/monoid.py b/effectful/ops/monoid.py
index 23791d29..e52d6bff 100644
--- a/effectful/ops/monoid.py
+++ b/effectful/ops/monoid.py
@@ -203,65 +203,12 @@ def __init__(self, name: str, identity: T, zero: T):
         self.zero = zero
 
 
+@dataclass(frozen=True)
 class Optimum[T]:
     """A value together with an assignment that attains it."""
 
-    __slots__ = ("assignment", "value")
-
-    def __init__(self, value: T, assignment: Mapping[Operation, Any] | None):
-        self.value = value
-        self.assignment = assignment
-
-    @property
-    def feasible(self) -> bool:
-        return self.assignment is not None
-
-    def __eq__(self, other: object) -> bool:
-        return (
-            isinstance(other, Optimum)
-            and self.value == other.value
-            and self.assignment == other.assignment
-        )
-
-    def __repr__(self) -> str:
-        return f"Optimum(value={self.value!r}, assignment={self.assignment!r})"
-
-
-@evaluate.register(Optimum)
-def _evaluate_optimum(expr: Optimum, **kwargs) -> Optimum:
-    """Evaluate an optimum's score and assignment values, preserving variable keys."""
-
-    assignment = (
-        None
-        if expr.assignment is None
-        else {variable: evaluate(value) for variable, value in expr.assignment.items()}
-    )
-    return Optimum(evaluate(expr.value), assignment)
-
-
-class ProductMonoid[L, R](Monoid[tuple[L, R]]):
-    """The componentwise product of two monoids.
-
-    ``ProductMonoid(left, right).plus`` combines the first components with
-    ``left.plus`` and the second components with ``right.plus``.
-    """
-
-    left: Monoid[L]
-    right: Monoid[R]
-
-    def __init__(self, left: Monoid[L], right: Monoid[R], name: str | None = None):
-        self.left = left
-        self.right = right
-        super().__init__(
-            name=name or f"{left.__name__}×{right.__name__}",
-            identity=(left.identity, right.identity),
-        )
-
-        # Product operations inherit these algebraic properties componentwise.
-        if is_commutative(left) and is_commutative(right):
-            is_commutative.register(self)
-        if is_idempotent(left) and is_idempotent(right):
-            is_idempotent.register(self)
+    value: T
+    assignment: Mapping[Operation, Any]
 
 
 Min = Monoid(name="Min", identity=float("inf"))
@@ -1634,20 +1581,6 @@ def plus(self, *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 MinOptimumPlus(ObjectInterpretation):
     """Lift :data:`Min` to values carrying minimizing assignments."""
 
@@ -1684,20 +1617,6 @@ def plus(self, *args):
             return fwd()
 
 
-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])
-
-
 def _disjoint_merge[K, V](*dicts: Mapping[K, V]) -> Mapping[K, V]:
     merged = {}
     for d in dicts:
@@ -1710,6 +1629,44 @@ def _disjoint_merge[K, V](*dicts: Mapping[K, V]) -> Mapping[K, V]:
     return merged
 
 
+class OptimumPlus(ObjectInterpretation):
+    @staticmethod
+    def _optimum_min_max(func, *args):
+        if (
+            not args
+            or not any(isinstance(a, Optimum) for a in args)
+            or any(isinstance(a, Optimum) and isinstance(a.value, Term) for a in args)
+        ):
+            return fwd()
+
+        return func(args, key=lambda a: a.value if isinstance(a, Optimum) else a)
+
+    @implements(Min.plus)
+    def _min_plus(self, *args):
+        return self._optimum_min_max(min, *args)
+
+    @implements(Max.plus)
+    def _max_plus(self, *args):
+        return self._optimum_min_max(max, *args)
+
+    @implements(Sum.plus)
+    def _sum_plus(self, *args):
+        if not args or not any(isinstance(arg, Optimum) for arg in args):
+            return fwd()
+        return Optimum(
+            Sum.plus(*(arg.value if isinstance(arg, Optimum) else arg for arg in args)),
+            _disjoint_merge(
+                *(arg.assignment if isinstance(arg, Optimum) else {} for arg in args)
+            ),
+        )
+
+    @implements(Monoid.plus)
+    def plus(self, monoid, *args):
+        if not args or not any(isinstance(arg, Optimum) for arg in args):
+            return fwd()
+        return monoid.plus(*(a.value if isinstance(a, Optimum) else a for a in args))
+
+
 class AssignmentPlus(ObjectInterpretation):
     """Disjoint-union implementation of :data:`Assignment`."""
 
@@ -1724,25 +1681,6 @@ def plus(self, *args):
         return _disjoint_merge(*args)
 
 
-class ProductMonoidPlus(ObjectInterpretation):
-    """Componentwise implementation of :class:`ProductMonoid`."""
-
-    @implements(Monoid.plus)
-    def plus(self, monoid, *args):
-        if not isinstance(monoid, ProductMonoid):
-            return fwd()
-        if not args:
-            return monoid.identity
-        if any(isinstance(arg, Term) for arg in args):
-            return fwd()
-        if not all(isinstance(arg, tuple) and len(arg) == 2 for arg in args):
-            return fwd()
-        return (
-            monoid.left.plus(*(arg[0] for arg in args)),
-            monoid.right.plus(*(arg[1] for arg in args)),
-        )
-
-
 class CartesianProductPlus(ObjectInterpretation):
     """Pure-Python implementation of :data:`CartesianProduct`."""
 
@@ -1849,35 +1787,6 @@ def plus(self, monoid, *args):
         return fwd()
 
 
-class SumOptimumPlus(ObjectInterpretation):
-    """Lift :data:`Sum` to values carrying minimizing assignments."""
-
-    @implements(Sum.plus)
-    def plus(self, *args):
-        if not args or not all(isinstance(arg, Optimum) for arg in args):
-            return fwd()
-        if any(not arg.feasible for arg in args):
-            return Optimum(float("inf"), None)
-        return Optimum(
-            Sum.plus(*(arg.value for arg in args)),
-            _disjoint_merge(*(typing.cast(Mapping, arg.assignment) for arg in args)),
-        )
-
-
-class PlusCastOptimum(ObjectInterpretation):
-    """Promote plain scores when adding an assignment-carrying weight."""
-
-    @implements(Sum.plus)
-    def plus(self, *args):
-        if any(isinstance(arg, Optimum) for arg in args) and any(
-            not isinstance(arg, Optimum) for arg in args
-        ):
-            return Sum.plus(
-                *(arg if isinstance(arg, Optimum) else Optimum(arg, {}) for arg in args)
-            )
-        return fwd()
-
-
 class EliminateSingletonStreams(ObjectInterpretation):
     """Eliminate a length-1 stream by substituting its sole element.
 
@@ -2289,20 +2198,13 @@ def extend(self, *intps: Interpretation) -> typing.Self:
     ReducePartial(),
     DeltaConcrete(),
     SumPlus(),
-    SumOptimumPlus(),
-    PlusCastOptimum(),
     MinPlus(),
-    MinOptimumPlus(),
     MaxPlus(),
-    MaxOptimumPlus(),
     ProductPlus(),
-    AssignmentPlus(),
-    ProductMonoidPlus(),
-    ArgMinPlus(),
-    ArgMaxPlus(),
     CartesianProductPlus(),
     UnionPlus(),
     IntersectionPlus(),
+    OptimumPlus(),
     ReduceWhereToMasks(),
 )
 
@@ -2354,7 +2256,7 @@ def extend(self, *intps: Interpretation) -> typing.Self:
     PlusOrder(),
     PlusCastFloat(),
     PlusCastIterable(),
-    SumOptimumPlus(),
+    OptimumPlus(),
     PlusCastOptimum(),
     MaskFusion(),
     MaskBool(),
diff --git a/tests/test_ops_monoid.py b/tests/test_ops_monoid.py
index b4f9f887..73db356c 100644
--- a/tests/test_ops_monoid.py
+++ b/tests/test_ops_monoid.py
@@ -41,7 +41,6 @@
     PlusPartial,
     PlusSingle,
     Product,
-    ProductMonoid,
     ReduceDependentRangeMask,
     ReduceDisjunctiveDisequalityMask,
     ReduceDistributeCartesianProduct,
@@ -582,25 +581,6 @@ def test_assignment_plus_rejects_duplicate_keys():
         Assignment.plus({"x": 1}, {"x": 1})
 
 
-def test_product_monoid_plus():
-    scored_assignment = ProductMonoid(Sum, Assignment)
-
-    with handler(EvaluateIntp):
-        result = scored_assignment.plus((2, {"x": 1}), (3, {"y": 2}))
-
-    assert result == (5, {"x": 1, "y": 2})
-    assert scored_assignment.identity == (0, {})
-    assert is_commutative(scored_assignment)
-    assert not is_idempotent(scored_assignment)
-
-
-def test_product_monoid_inherits_idempotence():
-    extrema = ProductMonoid(Min, Max)
-
-    assert is_commutative(extrema)
-    assert is_idempotent(extrema)
-
-
 def test_plus_partial():
     backend = IntBackend()
     x = backend.define_vars("x", ret="scalar")
@@ -1637,45 +1617,45 @@ def test_reduce_unfactor_reduces(Sum, Product, backend: Backend):
 
 
 def test_reduce_argmin(backend: Backend):
-    x, v = backend.define_vars("x", "v", ret="scalar")
+    x, xx, v = backend.define_vars("x", "xx", "v", ret="scalar")
 
     expr = Min.reduce(
         (x() - 1) ** 2,
-        {x: Sum.weighted(range(3), deffn(Optimum(Sum.identity, {x: v()}), v))},
+        {x: Sum.weighted(range(3), deffn(Optimum(Sum.identity, {xx: v()}), v))},
     )
 
     with handler(NormalizeIntp), handler(EvaluateIntp):
         result = evaluate(expr)
 
-    assert result == Optimum(0, {x: 1})
+    assert result == Optimum(0, {xx: 1})
 
 
 def test_reduce_argmin_sum(backend: Backend):
-    x, v = backend.define_vars("x", "v", ret="scalar")
+    x, xx, v = backend.define_vars("x", "xx", "v", ret="scalar")
 
     expr = Min.reduce(
         Sum.plus((x() - 1) ** 2, 2 * (x() - 2) ** 2),
-        {x: Sum.weighted(range(3), deffn(Optimum(Sum.identity, {x: v()}), v))},
+        {x: Sum.weighted(range(3), deffn(Optimum(Sum.identity, {xx: v()}), v))},
     )
 
     with handler(NormalizeIntp), handler(EvaluateIntp):
         result = evaluate(expr)
 
-    assert result == Optimum(1, {x: 2})
+    assert result == Optimum(1, {xx: 2})
 
 
 def test_reduce_argmin_sum_disjoint(backend: Backend):
-    x, y, v = backend.define_vars("x", "y", "v", ret="scalar")
+    x, xx, y, yy, v = backend.define_vars("x", "xx", "y", "yy", "v", ret="scalar")
 
     expr = Min.reduce(
         Sum.plus((x() - 1) ** 2, (y() - 2) ** 2),
         {
-            x: Sum.weighted(range(3), deffn(Optimum(Sum.identity, {x: v()}), v)),
-            y: Sum.weighted(range(3), deffn(Optimum(Sum.identity, {y: v()}), v)),
+            x: Sum.weighted(range(3), deffn(Optimum(Sum.identity, {xx: v()}), v)),
+            y: Sum.weighted(range(3), deffn(Optimum(Sum.identity, {yy: v()}), v)),
         },
     )
 
     with handler(NormalizeIntp), handler(EvaluateIntp):
         result = evaluate(expr)
 
-    assert result == Optimum(0, {x: 1, y: 2})
+    assert result == Optimum(0, {xx: 1, yy: 2})

From 58c3856ec63c0b90c76606fd72637da7d6d05234 Mon Sep 17 00:00:00 2001
From: Jack Feser <jack.feser@gmail.com>
Date: Wed, 29 Jul 2026 17:51:06 -0400
Subject: [PATCH 06/17] lint

---
 effectful/handlers/jax/monoid.py | 2 +-
 tests/test_ops_monoid.py         | 1 -
 2 files changed, 1 insertion(+), 2 deletions(-)

diff --git a/effectful/handlers/jax/monoid.py b/effectful/handlers/jax/monoid.py
index 20433e3d..f4c46285 100644
--- a/effectful/handlers/jax/monoid.py
+++ b/effectful/handlers/jax/monoid.py
@@ -294,7 +294,7 @@ def reduce(self, monoid: Monoid, body: Body, streams: Streams):
                 return fwd()
             assignment_vars[key] = value.op
 
-        if any(len(stream) == 0 for stream in streams.values()):
+        if any(len(typing.cast(range, stream)) == 0 for stream in streams.values()):
             return monoid.identity
 
         score_fvs = fvsof(body.value)
diff --git a/tests/test_ops_monoid.py b/tests/test_ops_monoid.py
index 73db356c..6e3c347a 100644
--- a/tests/test_ops_monoid.py
+++ b/tests/test_ops_monoid.py
@@ -63,7 +63,6 @@
     as_iterable,
     distributes_over,
     is_commutative,
-    is_idempotent,
     solve_group_equality,
 )
 from effectful.ops.semantics import coproduct, evaluate, fvsof, handler

From 1117fadd2d1496b281f182f190a964de8932bea2 Mon Sep 17 00:00:00 2001
From: Jack Feser <jack.feser@gmail.com>
Date: Wed, 29 Jul 2026 17:54:09 -0400
Subject: [PATCH 07/17] drop code

---
 effectful/ops/monoid.py | 16 ----------------
 1 file changed, 16 deletions(-)

diff --git a/effectful/ops/monoid.py b/effectful/ops/monoid.py
index e52d6bff..ea511077 100644
--- a/effectful/ops/monoid.py
+++ b/effectful/ops/monoid.py
@@ -213,8 +213,6 @@ class Optimum[T]:
 
 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 = Group(name="Sum", identity=0)
 Product = MonoidWithZero(name="Product", identity=1, zero=0)
 LogSumExp = Monoid(name="LogSumExp", identity=float("-inf"))
@@ -1667,20 +1665,6 @@ def plus(self, monoid, *args):
         return monoid.plus(*(a.value if isinstance(a, Optimum) else a for a in args))
 
 
-class AssignmentPlus(ObjectInterpretation):
-    """Disjoint-union implementation of :data:`Assignment`."""
-
-    @implements(Assignment.plus)
-    def plus(self, *args):
-        if not args:
-            return Assignment.identity
-        if any(isinstance(arg, Term) for arg in args):
-            return fwd()
-        if not all(isinstance(arg, Mapping) for arg in args):
-            return fwd()
-        return _disjoint_merge(*args)
-
-
 class CartesianProductPlus(ObjectInterpretation):
     """Pure-Python implementation of :data:`CartesianProduct`."""
 

From 65c4a66622564af3ed8896ba9a1ba7f34ed67d99 Mon Sep 17 00:00:00 2001
From: Jack Feser <jack.feser@gmail.com>
Date: Wed, 29 Jul 2026 18:04:10 -0400
Subject: [PATCH 08/17] remove bad tests

---
 tests/test_handlers_jax_monoid.py | 41 -------------------------------
 1 file changed, 41 deletions(-)

diff --git a/tests/test_handlers_jax_monoid.py b/tests/test_handlers_jax_monoid.py
index dabe4c27..ae7e54a1 100644
--- a/tests/test_handlers_jax_monoid.py
+++ b/tests/test_handlers_jax_monoid.py
@@ -128,47 +128,6 @@ def test_reduce_array_2(monoid, reductor, backend: JaxBackend):
     assert jnp.allclose(actual, expected)
 
 
-@pytest.mark.parametrize(
-    "monoid,expected_value,expected_assignment",
-    [(Min, 0, 1), (Max, 4, 3)],
-)
-def test_reduce_optimum_weighted_stream(
-    monoid, expected_value, expected_assignment, backend: JaxBackend
-):
-    x, v = backend.define_vars("x", "v", ret="scalar")
-    expr = monoid.reduce(
-        (x() - 1) ** 2,
-        {x: Sum.weighted(range(4), deffn(Optimum(0, {x: v()}), v))},
-    )
-
-    with handler(NormalizeIntp), handler(EvaluateIntp):
-        actual = evaluate(expr)
-
-    assert isinstance(actual, Optimum)
-    assert jnp.array_equal(actual.value, expected_value)
-    assert actual.assignment is not None
-    assert jnp.array_equal(actual.assignment[x], expected_assignment)
-
-
-def test_reduce_optimum_jit(backend: JaxBackend):
-    x = backend.define_vars("x", ret="scalar")
-
-    @jax.jit
-    def run(scores):
-        with handler(NormalizeIntp), handler(EvaluateIntp):
-            return Min.reduce(
-                Optimum(unbind_dims(scores, x), {x: x()}),
-                {x: range(scores.shape[0])},
-            )
-
-    actual = run(jnp.asarray([3.0, 1.0, 2.0]))
-
-    assert isinstance(actual, Optimum)
-    assert jnp.array_equal(actual.value, 1.0)
-    assert actual.assignment is not None
-    assert jnp.array_equal(actual.assignment[x], 1)
-
-
 SCALAR_PLUS = [
     pytest.param(Sum, 3.0, id="Sum"),
     pytest.param(Product, 2.0, id="Product"),

From 6c8e6f5267bab6d08f4e09ae0d3c922471c6fd38 Mon Sep 17 00:00:00 2001
From: Jack Feser <jack.feser@gmail.com>
Date: Wed, 29 Jul 2026 18:04:18 -0400
Subject: [PATCH 09/17] test that argmin is compatible with factorization

---
 tests/test_ops_monoid.py | 92 ++++++++++++++++++++++++++++++++++------
 1 file changed, 80 insertions(+), 12 deletions(-)

diff --git a/tests/test_ops_monoid.py b/tests/test_ops_monoid.py
index 6e3c347a..2edb94a7 100644
--- a/tests/test_ops_monoid.py
+++ b/tests/test_ops_monoid.py
@@ -12,7 +12,6 @@
 import effectful.handlers.jax.monoid  # noqa: F401
 from effectful.ops.monoid import (
     And,
-    Assignment,
     CartesianProduct,
     CartesianProductPlus,
     EliminateSingletonStreams,
@@ -569,17 +568,6 @@ def test_plus_zero(monoid, backend: Backend):
     backend.check_rewrite(lhs=lhs_left, rhs=rhs, rule={})
 
 
-def test_assignment_plus_disjoint_merge():
-    with handler(EvaluateIntp):
-        assert Assignment.plus({"x": 1}, {"y": 2}) == {"x": 1, "y": 2}
-        assert Assignment.plus() == {}
-
-
-def test_assignment_plus_rejects_duplicate_keys():
-    with handler(EvaluateIntp), pytest.raises(ValueError, match="Duplicate key"):
-        Assignment.plus({"x": 1}, {"x": 1})
-
-
 def test_plus_partial():
     backend = IntBackend()
     x = backend.define_vars("x", ret="scalar")
@@ -1361,6 +1349,86 @@ def test_reduce_weighted_factorization(backend: Backend):
     )
 
 
+def test_reduce_argmin_weighted_factorization(backend: Backend):
+    """Factoring a separable argmin preserves both minimizing assignments."""
+    x, xx, y, yy, v = backend.define_vars("x", "xx", "y", "yy", "v", ret="scalar")
+
+    lhs = Min.reduce(
+        Sum.plus(Optimum((x() - 1) ** 2, {}), Optimum((y() - 2) ** 2, {})),
+        {
+            x: Sum.weighted(range(3), deffn(Optimum(Sum.identity, {xx: v()}), v)),
+            y: Sum.weighted(range(4), deffn(Optimum(Sum.identity, {yy: v()}), v)),
+        },
+    )
+    rhs = Sum.plus(
+        Min.reduce(
+            Sum.plus(
+                Optimum(Sum.identity, {xx: x()}),
+                Sum.plus(Optimum((x() - 1) ** 2, {})),
+            ),
+            {x: range(3)},
+        ),
+        Min.reduce(
+            Sum.plus(
+                Optimum(Sum.identity, {yy: y()}),
+                Sum.plus(Optimum((y() - 2) ** 2, {})),
+            ),
+            {y: range(4)},
+        ),
+    )
+
+    backend.check_rewrite(
+        lhs=lhs, rhs=rhs, rule=coproduct(ReduceWeightedStream(), Factor())
+    )
+
+
+def test_reduce_argmin_weighted_repeated_factorization(backend: Backend):
+    """Factoring repeatedly preserves every weighted minimizing assignment."""
+    x, xx, y, yy, z, zz, v = backend.define_vars(
+        "x", "xx", "y", "yy", "z", "zz", "v", ret="scalar"
+    )
+
+    lhs = Min.reduce(
+        Sum.plus(
+            Optimum((x() - 1) ** 2, {}),
+            Optimum((y() - 2) ** 2, {}),
+            Optimum((z() - 3) ** 2, {}),
+        ),
+        {
+            x: Sum.weighted(range(3), deffn(Optimum(Sum.identity, {xx: v()}), v)),
+            y: Sum.weighted(range(4), deffn(Optimum(Sum.identity, {yy: v()}), v)),
+            z: Sum.weighted(range(5), deffn(Optimum(Sum.identity, {zz: v()}), v)),
+        },
+    )
+    rhs = Sum.plus(
+        Min.reduce(
+            Sum.plus(
+                Optimum(Sum.identity, {xx: x()}),
+                Sum.plus(Optimum((x() - 1) ** 2, {})),
+            ),
+            {x: range(3)},
+        ),
+        Min.reduce(
+            Sum.plus(
+                Optimum(Sum.identity, {yy: y()}),
+                Sum.plus(Optimum((y() - 2) ** 2, {})),
+            ),
+            {y: range(4)},
+        ),
+        Min.reduce(
+            Sum.plus(
+                Optimum(Sum.identity, {zz: z()}),
+                Sum.plus(Optimum((z() - 3) ** 2, {})),
+            ),
+            {z: range(5)},
+        ),
+    )
+
+    backend.check_rewrite(
+        lhs=lhs, rhs=rhs, rule=coproduct(ReduceWeightedStream(), Factor())
+    )
+
+
 def test_weighted_expectation_demo():
     """Demo: compute E[f(X)] = Σ_x w(x)·f(x) via a weighted reduce.
 

From d440a01b502ca9a2ba58555f457dcc66e4f3f5b1 Mon Sep 17 00:00:00 2001
From: Jack Feser <jack.feser@gmail.com>
Date: Thu, 30 Jul 2026 10:01:58 -0400
Subject: [PATCH 10/17] format

---
 tests/test_handlers_jax_monoid.py | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/tests/test_handlers_jax_monoid.py b/tests/test_handlers_jax_monoid.py
index ae7e54a1..a538664d 100644
--- a/tests/test_handlers_jax_monoid.py
+++ b/tests/test_handlers_jax_monoid.py
@@ -21,15 +21,11 @@
 from effectful.ops.monoid import (
     EliminateSingletonStreams,
     EvaluateIntp,
-    Max,
-    Min,
     NormalizeIntp,
-    Optimum,
     Product,
     Sum,
 )
 from effectful.ops.semantics import coproduct, evaluate, handler
-from effectful.ops.syntax import deffn
 from tests._monoid_helpers import JaxBackend
 
 MONOIDS = [

From 930ee20491a712a2c22257d054950917cb300a61 Mon Sep 17 00:00:00 2001
From: Jack Feser <jack.feser@gmail.com>
Date: Mon, 10 Aug 2026 18:59:26 -0400
Subject: [PATCH 11/17] add gradient-descent optimization over continuous
 spaces

---
 effectful/handlers/numpyro.py  | 137 ++++++++++++++++++++++++++++++++-
 tests/test_handlers_numpyro.py |  88 ++++++++++++++++++++-
 2 files changed, 221 insertions(+), 4 deletions(-)

diff --git a/effectful/handlers/numpyro.py b/effectful/handlers/numpyro.py
index 818f5b3b..dcd95a6f 100644
--- a/effectful/handlers/numpyro.py
+++ b/effectful/handlers/numpyro.py
@@ -1,6 +1,7 @@
 try:
     import numpyro
     import numpyro.distributions as dist
+    import numpyro.optim
 except ImportError:
     raise ImportError("Numpyro is required to use effectful.handlers.numpyro")
 
@@ -16,14 +17,16 @@
 from effectful.handlers.jax._handlers import _register_jax_op, is_eager_array
 from effectful.ops.monoid import (
     LogSumExp,
+    Min,
     Monoid,
     NormalizeIntp,
+    Optimum,
     Product,
     Stream,
     Streams,
     Sum,
 )
-from effectful.ops.semantics import evaluate, fwd, typeof
+from effectful.ops.semantics import evaluate, fvsof, fwd, handler, typeof
 from effectful.ops.syntax import ObjectInterpretation, defdata, deffn, defop, implements
 from effectful.ops.types import NotHandled, Operation, Term
 
@@ -1194,7 +1197,6 @@ def _embed_independent(d: dist.Independent) -> Term[dist.Independent]:
     return Independent(d.base_dist, d.reinterpreted_batch_ndims)
 
 
-@Operation.define
 def distribution_stream(
     distribution: numpyro.distributions.Distribution,
 ) -> Stream[jax.Array]:
@@ -1237,4 +1239,135 @@ def _(self, monoid, body, streams: Streams):
         return fwd()
 
 
+@Operation.define
+def constraint_stream(
+    constraint: numpyro.distributions.constraints.Constraint, prototype: jax.Array
+) -> Stream[jax.Array]:
+    raise NotHandled
+
+
+class NestConstraintMinReduce(ObjectInterpretation):
+    """Move constraint streams into an inner :data:`~effectful.ops.monoid.Min`.
+
+    This rewrites::
+
+        Min.reduce(body, constraint_streams | other_streams)
+
+    to::
+
+        Min.reduce(Min.reduce(body, constraint_streams), other_streams)
+
+    when both partitions are nonempty. A constraint stream may depend on an
+    outer stream, but an outer stream may not depend on a constraint variable;
+    the latter ordering would move the variable outside its scope and is left
+    for another handler.
+    """
+
+    @implements(Min.reduce)
+    def reduce(self, body, streams):
+        constraint_streams = {
+            key: stream
+            for key, stream in streams.items()
+            if isinstance(stream, Term) and stream.op is constraint_stream
+        }
+        if not constraint_streams or len(constraint_streams) == len(streams):
+            return fwd()
+
+        other_streams = {
+            key: stream
+            for key, stream in streams.items()
+            if key not in constraint_streams
+        }
+        if fvsof(other_streams) & set(constraint_streams):
+            return fwd()
+
+        return Min.reduce(Min.reduce(body, constraint_streams), other_streams)
+
+
+class AdamConstraintMinReduce(ObjectInterpretation):
+    """Minimize an all-continuous constraint-stream bundle with Adam.
+
+    Optimization takes place in unconstrained coordinates using NumPyro's
+    ``biject_to`` transforms. The prototypes carried by
+    :func:`constraint_stream` determine shape and dtype; each constraint's
+    :meth:`~numpyro.distributions.constraints.Constraint.feasible_like` method
+    constructs the constrained initial value.
+    """
+
+    def __init__(
+        self,
+        step_size=1e-2,
+        *,
+        num_steps: int = 1_000,
+        b1: float = 0.9,
+        b2: float = 0.999,
+        eps: float = 1e-8,
+    ):
+        if num_steps < 0:
+            raise ValueError("num_steps must be nonnegative")
+        self.num_steps = num_steps
+        self.optimizer = numpyro.optim.Adam(step_size=step_size, b1=b1, b2=b2, eps=eps)
+
+    @implements(Min.reduce)
+    def reduce(self, body, streams):
+        if not streams or not all(
+            isinstance(stream, Term) and stream.op is constraint_stream
+            for stream in streams.values()
+        ):
+            return fwd()
+
+        constraints = tuple(stream.args[0] for stream in streams.values())
+        if any(constraint.is_discrete for constraint in constraints):
+            return fwd()
+
+        score = body.value if isinstance(body, Optimum) else body
+        stream_keys = tuple(streams)
+
+        transforms = tuple(
+            numpyro.distributions.transforms.biject_to(constraint)
+            for constraint in constraints
+        )
+        prototypes = tuple(stream.args[1] for stream in streams.values())
+        constrained_initial = tuple(
+            constraint.feasible_like(prototype)
+            for constraint, prototype in zip(constraints, prototypes, strict=True)
+        )
+        unconstrained_initial = tuple(
+            transform.inv(value)
+            for transform, value in zip(transforms, constrained_initial, strict=True)
+        )
+
+        def substitute(value, unconstrained):
+            constrained = tuple(
+                transform(x)
+                for transform, x in zip(transforms, unconstrained, strict=True)
+            )
+            substitutions = {
+                key: (lambda value=value: value)
+                for key, value in zip(stream_keys, constrained, strict=True)
+            }
+            with handler(substitutions):
+                return evaluate(value)
+
+        def objective(unconstrained):
+            return substitute(score, unconstrained)
+
+        initial_score = objective(unconstrained_initial)
+        if isinstance(initial_score, Term):
+            return fwd()
+        if jnp.ndim(initial_score) != 0:
+            raise ValueError("Min objective must be scalar")
+
+        state = self.optimizer.init(unconstrained_initial)
+
+        def step(_, state):
+            (_, _), state = self.optimizer.eval_and_stable_update(
+                lambda params: (objective(params), None), state
+            )
+            return state
+
+        state = jax.lax.fori_loop(0, self.num_steps, step, state)
+        return substitute(body, self.optimizer.get_params(state))
+
+
 NormalizeIntp.extend(ReduceEnumerableDistribution())
diff --git a/tests/test_handlers_numpyro.py b/tests/test_handlers_numpyro.py
index 8a9b03b3..9dab2d68 100644
--- a/tests/test_handlers_numpyro.py
+++ b/tests/test_handlers_numpyro.py
@@ -12,8 +12,8 @@
 import effectful.handlers.jax.numpy as jnp
 import effectful.handlers.numpyro as dist
 from effectful.handlers.jax import bind_dims, jax_getitem, sizesof, unbind_dims
-from effectful.ops.monoid import LogSumExp, Product, Sum
-from effectful.ops.semantics import typeof
+from effectful.ops.monoid import LogSumExp, Min, Optimum, Product, Sum
+from effectful.ops.semantics import handler, typeof
 from effectful.ops.syntax import deffn, defop
 from effectful.ops.types import Operation, Term
 from tests._monoid_helpers import JaxBackend
@@ -1058,3 +1058,87 @@ def model():
         mcmc.run(jr.PRNGKey(0))
 
     assert mcmc.get_samples()["theta"].shape == (20, 3)
+
+
+def test_nest_constraint_min_reduce():
+    x = defop(jax.Array, name="x")
+    y = defop(jax.Array, name="y")
+    constrained = dist.constraint_stream(
+        numpyro.distributions.constraints.positive, jnp.asarray(1.0)
+    )
+
+    with handler(dist.NestConstraintMinReduce()):
+        result = Min.reduce(x(), {x: constrained, y: range(3)})
+
+    assert isinstance(result, Term) and result.op is Min.reduce
+    inner, outer_streams = result.args
+    assert list(outer_streams.values()) == [range(3)]
+    assert isinstance(inner, Term) and inner.op is Min.reduce
+    assert len(inner.args[1]) == 1
+    inner_stream = next(iter(inner.args[1].values()))
+    assert isinstance(inner_stream, Term) and inner_stream.op is dist.constraint_stream
+
+
+def test_adam_constraint_min_reduce():
+    x = defop(jax.Array, name="x")
+    y = defop(jax.Array, name="y")
+    streams = {
+        x: dist.constraint_stream(
+            numpyro.distributions.constraints.positive, jnp.asarray(0.5)
+        ),
+        y: dist.constraint_stream(
+            numpyro.distributions.constraints.real, jnp.asarray(0.0)
+        ),
+    }
+    objective = Optimum(
+        (x() - 2.0) ** 2 + (y() + 1.0) ** 2,
+        {"x": x(), "y": y()},
+    )
+
+    with handler(dist.AdamConstraintMinReduce(step_size=0.05, num_steps=500)):
+        result = Min.reduce(objective, streams)
+
+    assert isinstance(result, Optimum)
+    assert jnp.allclose(result.value, 0.0, atol=1e-5)
+    assert jnp.allclose(result.assignment["x"], 2.0, atol=1e-4)
+    assert jnp.allclose(result.assignment["y"], -1.0, atol=1e-4)
+
+
+def test_adam_constraint_min_reduce_initializes_with_feasible_like():
+    x = defop(jax.Array, name="x")
+    constraint = numpyro.distributions.constraints.interval(-2.0, 4.0)
+    stream = dist.constraint_stream(constraint, jnp.asarray(100.0))
+
+    with handler(dist.AdamConstraintMinReduce(num_steps=0)):
+        result = Min.reduce(x(), {x: stream})
+
+    assert jnp.allclose(result, 1.0)
+
+
+def test_adam_constraint_min_reduce_forwards_mixed_bundle():
+    x = defop(jax.Array, name="x")
+    y = defop(jax.Array, name="y")
+    constrained = dist.constraint_stream(
+        numpyro.distributions.constraints.real, jnp.asarray(0.0)
+    )
+
+    with handler(dist.AdamConstraintMinReduce()):
+        result = Min.reduce(x() ** 2, {x: constrained, y: range(2)})
+
+    assert isinstance(result, Term) and result.op is Min.reduce
+    assert len(result.args[1]) == 2
+
+
+def test_nest_constraint_min_reduce_preserves_stream_dependencies():
+    x = defop(jax.Array, name="x")
+    y = defop(jax.Array, name="y")
+    constrained = dist.constraint_stream(
+        numpyro.distributions.constraints.positive, jnp.asarray(1.0)
+    )
+
+    with handler(dist.NestConstraintMinReduce()):
+        result = Min.reduce(x(), {x: constrained, y: (x(),)})
+
+    assert isinstance(result, Term) and result.op is Min.reduce
+    assert len(result.args[1]) == 2
+    assert not (isinstance(result.args[0], Term) and result.args[0].op is Min.reduce)

From b99529be04a9626bcd703ba08b2962825f6b7673 Mon Sep 17 00:00:00 2001
From: Jack Feser <jack.feser@gmail.com>
Date: Tue, 11 Aug 2026 17:35:57 -0400
Subject: [PATCH 12/17] wip

---
 effectful/ops/monoid.py | 106 +++++++++++++++++++---------------------
 1 file changed, 49 insertions(+), 57 deletions(-)

diff --git a/effectful/ops/monoid.py b/effectful/ops/monoid.py
index ea511077..8dc6abd7 100644
--- a/effectful/ops/monoid.py
+++ b/effectful/ops/monoid.py
@@ -1579,42 +1579,6 @@ def plus(self, *args):
         )
 
 
-class MinOptimumPlus(ObjectInterpretation):
-    """Lift :data:`Min` to values carrying minimizing assignments."""
-
-    @implements(Min.plus)
-    def plus(self, *args):
-        if not args or not all(isinstance(a, Optimum) for a in args):
-            return fwd()
-        feasible = [a for a in args if a.feasible]
-        if not feasible:
-            return Optimum(Min.identity, None)
-        if any(isinstance(a.value, Term) for a in feasible):
-            return fwd()
-        try:
-            return min(feasible, key=lambda a: a.value)
-        except (TypeError, ValueError):
-            return fwd()
-
-
-class MaxOptimumPlus(ObjectInterpretation):
-    """Lift :data:`Max` to values carrying maximizing assignments."""
-
-    @implements(Max.plus)
-    def plus(self, *args):
-        if not args or not all(isinstance(a, Optimum) for a in args):
-            return fwd()
-        feasible = [a for a in args if a.feasible]
-        if not feasible:
-            return Optimum(Max.identity, None)
-        if any(isinstance(a.value, Term) for a in feasible):
-            return fwd()
-        try:
-            return max(feasible, key=lambda a: a.value)
-        except (TypeError, ValueError):
-            return fwd()
-
-
 def _disjoint_merge[K, V](*dicts: Mapping[K, V]) -> Mapping[K, V]:
     merged = {}
     for d in dicts:
@@ -1627,17 +1591,22 @@ def _disjoint_merge[K, V](*dicts: Mapping[K, V]) -> Mapping[K, V]:
     return merged
 
 
-class OptimumPlus(ObjectInterpretation):
-    @staticmethod
-    def _optimum_min_max(func, *args):
-        if (
-            not args
-            or not any(isinstance(a, Optimum) for a in args)
-            or any(isinstance(a, Optimum) and isinstance(a.value, Term) for a in args)
-        ):
-            return fwd()
+class PlusOptimum(ObjectInterpretation):
+    """Sum values that are annotated with a minimizing (or maximizing) assignment."""
+
+    def _is_reducible(self, args):
+        return (
+            args
+            and any(isinstance(a, Optimum) for a in args)
+            and all(not fvsof(a.value if isinstance(a, Optimum) else a) for a in args)
+        )
 
-        return func(args, key=lambda a: a.value if isinstance(a, Optimum) else a)
+    def _optimum_min_max(self, func, *args):
+        return (
+            func(args, key=lambda a: a.value if isinstance(a, Optimum) else a)
+            if self._is_reducible(args)
+            else fwd()
+        )
 
     @implements(Min.plus)
     def _min_plus(self, *args):
@@ -1649,20 +1618,43 @@ def _max_plus(self, *args):
 
     @implements(Sum.plus)
     def _sum_plus(self, *args):
-        if not args or not any(isinstance(arg, Optimum) for arg in args):
-            return fwd()
-        return Optimum(
-            Sum.plus(*(arg.value if isinstance(arg, Optimum) else arg for arg in args)),
-            _disjoint_merge(
-                *(arg.assignment if isinstance(arg, Optimum) else {} for arg in args)
-            ),
+        return (
+            Optimum(
+                Sum.plus(
+                    *(arg.value if isinstance(arg, Optimum) else arg for arg in args)
+                ),
+                _disjoint_merge(
+                    *(
+                        arg.assignment if isinstance(arg, Optimum) else {}
+                        for arg in args
+                    )
+                ),
+            )
+            if self._is_reducible(args)
+            else fwd()
         )
 
     @implements(Monoid.plus)
     def plus(self, monoid, *args):
-        if not args or not any(isinstance(arg, Optimum) for arg in args):
-            return fwd()
-        return monoid.plus(*(a.value if isinstance(a, Optimum) else a for a in args))
+        return (
+            monoid.plus(*(a.value if isinstance(a, Optimum) else a for a in args))
+            if self._is_reducible(args)
+            else fwd()
+        )
+
+
+class PlusCastOptimum(ObjectInterpretation):
+    """Upcast non-Optimum arguments to Monoid.plus."""
+
+    @implements(Monoid.plus)
+    def plus(self, monoid, *args):
+        num_optimum = sum(isinstance(a, Optimum) for a in args)
+        if 0 < num_optimum < len(args):
+            new_args = (
+                Optimum(a, {}) if not isinstance(a, Optimum) else a for a in args
+            )
+            return monoid.plus(*new_args)
+        return fwd()
 
 
 class CartesianProductPlus(ObjectInterpretation):
@@ -2188,7 +2180,7 @@ def extend(self, *intps: Interpretation) -> typing.Self:
     CartesianProductPlus(),
     UnionPlus(),
     IntersectionPlus(),
-    OptimumPlus(),
+    PlusOptimum(),
     ReduceWhereToMasks(),
 )
 

From b8ce0e80232ed4a375ca44faa7dd2ac8a1eb1c30 Mon Sep 17 00:00:00 2001
From: Jack Feser <jack.feser@gmail.com>
Date: Tue, 11 Aug 2026 12:43:06 -0400
Subject: [PATCH 13/17] drop

---
 effectful/ops/monoid.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/effectful/ops/monoid.py b/effectful/ops/monoid.py
index 8dc6abd7..4a96e190 100644
--- a/effectful/ops/monoid.py
+++ b/effectful/ops/monoid.py
@@ -2246,8 +2246,6 @@ def extend(self, *intps: Interpretation) -> typing.Self:
     MinPlus(),
     MaxPlus(),
     ProductPlus(),
-    ArgMinPlus(),
-    ArgMaxPlus(),
     CartesianProductPlus(),
     UnionPlus(),
     AndPlus(),

From c83057f13e95f93c335dabffdf74725db99863e5 Mon Sep 17 00:00:00 2001
From: Jack Feser <jack.feser@gmail.com>
Date: Tue, 11 Aug 2026 12:43:47 -0400
Subject: [PATCH 14/17] rename

---
 effectful/ops/monoid.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/effectful/ops/monoid.py b/effectful/ops/monoid.py
index 4a96e190..2dd0c1e0 100644
--- a/effectful/ops/monoid.py
+++ b/effectful/ops/monoid.py
@@ -2232,7 +2232,7 @@ def extend(self, *intps: Interpretation) -> typing.Self:
     PlusOrder(),
     PlusCastFloat(),
     PlusCastIterable(),
-    OptimumPlus(),
+    PlusOptimum(),
     PlusCastOptimum(),
     MaskFusion(),
     MaskBool(),

From c16af6561c8db8e7af7288efdffbefaea2bb4773 Mon Sep 17 00:00:00 2001
From: Jack Feser <jack.feser@gmail.com>
Date: Tue, 11 Aug 2026 17:47:30 -0400
Subject: [PATCH 15/17] fix missing import

---
 tests/test_ops_monoid.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/test_ops_monoid.py b/tests/test_ops_monoid.py
index 2edb94a7..7cc97454 100644
--- a/tests/test_ops_monoid.py
+++ b/tests/test_ops_monoid.py
@@ -62,6 +62,7 @@
     as_iterable,
     distributes_over,
     is_commutative,
+    is_idempotent,
     solve_group_equality,
 )
 from effectful.ops.semantics import coproduct, evaluate, fvsof, handler

From 471c0db63f125691bfa44bf676ede640d95f74a5 Mon Sep 17 00:00:00 2001
From: Jack Feser <jack.feser@gmail.com>
Date: Tue, 11 Aug 2026 18:13:17 -0400
Subject: [PATCH 16/17] fix syntactic_hash bug

---
 effectful/ops/syntax.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/effectful/ops/syntax.py b/effectful/ops/syntax.py
index c41fcc9c..7b0d192c 100644
--- a/effectful/ops/syntax.py
+++ b/effectful/ops/syntax.py
@@ -1039,7 +1039,7 @@ def syntactic_hash(__dispatch: Callable[[type], Callable[[Any], int]], x) -> int
     :param x: A term.
     :returns: An integer hash.
     """
-    if dataclasses.is_dataclass(x) and not isinstance(x, type):
+    if dataclasses.is_dataclass(x) and not isinstance(x, type | Term):
         return hash(
             (
                 "dataclass",

From ad283b16f9289ad62e5d4fdd996c6d50046c1635 Mon Sep 17 00:00:00 2001
From: Jack Feser <jack.feser@gmail.com>
Date: Tue, 11 Aug 2026 18:25:08 -0400
Subject: [PATCH 17/17] fix test failures

---
 effectful/handlers/jax/monoid.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/effectful/handlers/jax/monoid.py b/effectful/handlers/jax/monoid.py
index f4c46285..fe528aa4 100644
--- a/effectful/handlers/jax/monoid.py
+++ b/effectful/handlers/jax/monoid.py
@@ -97,8 +97,10 @@ 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
+        if (
+            any(_is_jax(t) for t in arg_types)
+            and any(not _is_jax(t) for t in arg_types)
+            and all(issubclass(t, jax.typing.ArrayLike) for t in arg_types)
         ):
             return monoid.plus(
                 *(