Generalize ReduceEqualityMaskRange

effectful/handlers/jax/monoid.py

@@ -116,6 +116,14 @@ def inverse(self, value):
         return jnp.negative(value)
 
 
+class SumInverseJax(ObjectInterpretation):
+    @implements(Sum.inverse)
+    def inverse(self, value):
+        if not _jax_args((value,)):
+            return fwd()
+        return jnp.negative(value)
+
+
 class ProductPlusJax(ObjectInterpretation):
     @implements(Product.plus)
     def plus(self, *args):
@@ -801,6 +809,7 @@ def einsum(
 
 EvaluateIntp.extend(
     SumPlusJax(),
+    SumInverseJax(),
     ProductPlusJax(),
     MinPlusJax(),
     MaxPlusJax(),