hacks

Author: jfeser

effectful/handlers/jax/monoid.py

@@ -174,9 +174,16 @@ def reduce(self, monoid, body, streams):
         if typeof(body) is not jax.Array:
             return fwd()
 
-        # delta bodies belong to ReduceDeltaIndependent
-        if isinstance(body, Term) and (body.op is delta or _is_monoid_plus(body.op)):
-            return fwd()
+        if isinstance(body, Term):
+            # delta bodies belong to ReduceDeltaIndependent
+            if body.op is delta:
+                return fwd()
+            if _is_monoid_plus(body.op):
+                # factorizable bodies get factored
+                factors = body.args
+                factor_fvs = [fvsof(f) for f in factors]
+                if not all(k in fvs for fvs in factor_fvs for k in streams):
+                    return fwd()
 
         body_fvs = fvsof(body)
         used = [

effectful/ops/monoid.py

@@ -828,6 +828,22 @@ def reduce(self, monoid, body, streams):
         return result
 
 
+class PlusInf(ObjectInterpretation):
+    """Workaround for the inability to give Monoid.plus(x, inf) a type."""
+
+    @implements(Monoid.plus)
+    def plus(self, monoid, *args):
+        if monoid in (Sum, Max) and any(
+            not isinstance(x, Term) and x == float("inf") for x in args
+        ):
+            return float("inf")
+        if monoid in (Sum, Min) and any(
+            not isinstance(x, Term) and x == -float("inf") for x in args
+        ):
+            return -float("inf")
+        return fwd()
+
+
 class _ExtensibleInterpretation(UserDict, Interpretation):
     def extend(self, *intps: Interpretation) -> typing.Self:
         for intp in intps:
@@ -859,6 +875,7 @@ def extend(self, *intps: Interpretation) -> typing.Self:
     ArgMinPlus(),
     ArgMaxPlus(),
     CartesianProductPlus(),
+    PlusInf(),
 )
 """``NormalizeIntp``applies pure-Term rewrites (associativity, distributivity,
 identity elimination, fusion, factorization, etc.).

tests/test_handlers_jax_monoid.py

@@ -333,12 +333,10 @@ def test_reduce_contraction_single(backend: JaxBackend):
     )
 
     lhs = Sum.reduce(Product.plus(A(i()), B(i())), {i: I()})
-    rhs = Product.plus(
-        jnp.tensordot(
-            bind_dims(A(unbind_dims(I(), i)), i),
-            bind_dims(B(unbind_dims(I(), i)), i),
-            axes=((0,), (0,)),
-        )
+    rhs = jnp.tensordot(
+        bind_dims(A(unbind_dims(I(), i)), i),
+        bind_dims(B(unbind_dims(I(), i)), i),
+        axes=((0,), (0,)),
     )
     backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceSumProductContraction())
 
@@ -351,79 +349,10 @@ def test_reduce_contraction_double(backend: JaxBackend):
     )
 
     lhs = Sum.reduce(Product.plus(A(i(), j()), B(i(), j())), {i: I(), j: J()})
-    rhs = Product.plus(
-        jnp.tensordot(
-            bind_dims(A(unbind_dims(I(), i), unbind_dims(J(), j)), i, j),
-            bind_dims(B(unbind_dims(I(), i), unbind_dims(J(), j)), i, j),
-            axes=((0, 1), (0, 1)),
-        )
-    )
-    backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceSumProductContraction())
-
-
-def test_reduce_contraction_multi_factor(backend: JaxBackend):
-    i, j = backend.define_vars("i", "j", ret="scalar")
-    I, J = backend.define_vars("I", "J", ret="stream")
-    B = backend.define_vars(
-        "B", arg_types=(backend.scalar_typ, backend.scalar_typ), ret="scalar"
-    )
-    A, C = backend.define_vars("A", "C", arg_types=(backend.scalar_typ,), ret="scalar")
-
-    lhs = Sum.reduce(Product.plus(A(i()), B(i(), j()), C(j())), {i: I(), j: J()})
-    rhs = Product.plus(
-        jnp.tensordot(
-            bind_dims(
-                jnp.tensordot(
-                    bind_dims(A(unbind_dims(I(), i)), i),
-                    bind_dims(B(unbind_dims(I(), i), unbind_dims(J(), j)), i),
-                    axes=((0,), (0,)),
-                ),
-                j,
-            ),
-            bind_dims(C(unbind_dims(J(), j)), j),
-            axes=((0,), (0,)),
-        )
-    )
-    backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceSumProductContraction())
-
-
-def test_reduce_contraction_shared_only(backend: JaxBackend):
-    i, j, k = backend.define_vars("i", "j", "k", ret="scalar")
-    I, J, K = backend.define_vars("I", "J", "K", ret="stream")
-    A = backend.define_vars(
-        "A",
-        arg_types=(backend.scalar_typ, backend.scalar_typ, backend.scalar_typ),
-        ret="scalar",
-    )
-    B, C = backend.define_vars(
-        "B", "C", arg_types=(backend.scalar_typ, backend.scalar_typ), ret="scalar"
-    )
-
-    lhs = Sum.reduce(
-        Product.plus(A(i(), j(), k()), B(i(), j()), C(j(), k())),
-        {i: I(), j: J(), k: K()},
-    )
-    rhs = Product.plus(
-        jnp.tensordot(
-            bind_dims(
-                jnp.tensordot(
-                    bind_dims(
-                        A(
-                            unbind_dims(I(), i),
-                            unbind_dims(J(), j),
-                            unbind_dims(K(), k),
-                        ),
-                        i,
-                    ),
-                    bind_dims(B(unbind_dims(I(), i), unbind_dims(J(), j)), i),
-                    axes=((0,), (0,)),
-                ),
-                j,
-                k,
-            ),
-            bind_dims(C(unbind_dims(J(), j), unbind_dims(K(), k)), j, k),
-            axes=((0, 1), (0, 1)),
-        )
+    rhs = jnp.tensordot(
+        bind_dims(A(unbind_dims(I(), i), unbind_dims(J(), j)), i, j),
+        bind_dims(B(unbind_dims(I(), i), unbind_dims(J(), j)), i, j),
+        axes=((0, 1), (0, 1)),
     )
     backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceSumProductContraction())