add support for pymbolic.EqualityMapper

Author: alexfikl

pytato/scalar_expr.py

@@ -41,6 +41,8 @@
         DistributeMapperBase)
 from pymbolic.mapper.stringifier import (StringifyMapper as
         StringifyMapperBase)
+from pymbolic.mapper.equality import (EqualityMapper as
+        EqualityMapperBase)
 from pymbolic.mapper import CombineMapper as CombineMapperBase
 from pymbolic.mapper.collector import TermCollector as TermCollectorBase
 from immutables import Map
@@ -184,6 +186,20 @@ def map_reduce(self, expr: Any, enclosing_prec: Any, *args: Any) -> str:
         bounds_expr = "{" + bounds_expr + "}"
         return (f"{expr.op}({bounds_expr}, {self.rec(expr.inner_expr, PN)})")
 
+
+class EqualityMapper(EqualityMapperBase):
+    def map_reduce(self, expr: Reduce, other: Reduce) -> bool:
+        return (
+                len(expr.bounds) == len(other.bounds)
+                and all(k == other_k
+                        and self.rec(lb, other_lb) and self.rec(ub, other_ub)
+                    for (k, (lb, ub)), (other_k, (other_lb, other_ub)) in zip(
+                        sorted(expr.bounds.items()),
+                        sorted(other.bounds.items())))
+                and expr.op == other.op
+                and self.rec(expr.inner_expr, other.inner_expr)
+                )
+
 # }}}
 
 
@@ -240,6 +256,9 @@ def distribute(expr: Any, parameters: FrozenSet[Any] = frozenset(),
 # {{{ custom scalar expression nodes
 
 class ExpressionBase(prim.Expression):
+    def make_equality_mapper(self) -> EqualityMapper:
+        return EqualityMapper()
+
     def make_stringifier(self, originating_stringifier: Any = None) -> str:
         return StringifyMapper()
 

test/test_pytato.py

@@ -355,6 +355,8 @@ def test_userscollector():
 
 
 def test_asciidag():
+    pytest.importorskip("asciidag")
+
     n = pt.make_size_param("n")
     array = pt.make_placeholder(name="array", shape=n, dtype=np.float64)
     stack = pt.stack([array, 2*array, array + 6])