Implement IdentityMapper.map_reduce

inducer · Jul 1, 2024 · 5810fba · 5810fba
1 parent 71c7cc7
commit 5810fba
Showing 1 changed file with 11 additions and 1 deletion.
diff --git a/pytato/scalar_expr.py b/pytato/scalar_expr.py
@@ -111,7 +111,17 @@ def map_reduce(self, expr: Reduce, *args: Any, **kwargs: Any) -> Any:
 
 
 class IdentityMapper(IdentityMapperBase):
-    pass
+    def map_reduce(self, expr: Reduce, *args: Any, **kwargs: Any) -> Any:
+        return Reduce(
+                      self.rec(expr.inner_expr, *args, **kwargs),
+                      expr.op,
+                      immutabledict({
+                                        name: (
+                                            self.rec(lower, *args, **kwargs),
+                                            self.rec(upper, *args, **kwargs)
+                                        )
+                                        for name, (lower, upper) in expr.bounds.items()
+                                    }))
 
     def map_type_cast(self, expr: TypeCast, *args: Any, **kwargs: Any) -> Any:
         return TypeCast(expr.dtype, self.rec(expr.inner_expr, *args, **kwargs))