fix[next]: Fix type preservation in CSE

src/gt4py/next/iterator/transforms/cse.py

@@ -14,6 +14,7 @@
 import operator
 from typing import Callable, Iterable, TypeVar, Union, cast
 
+import gt4py.next.iterator.ir_utils.ir_makers as im
 from gt4py.eve import (
     NodeTranslator,
     NodeVisitor,
@@ -241,7 +242,6 @@ def extract_subexpression(
     Examples:
         Default case for `(x+y) + ((x+y)+z)`:
 
-        >>> import gt4py.next.iterator.ir_utils.ir_makers as im
         >>> from gt4py.eve.utils import UIDGenerator
         >>> expr = im.plus(im.plus("x", "y"), im.plus(im.plus("x", "y"), "z"))
         >>> predicate = lambda subexpr, num_occurences: num_occurences > 1
@@ -433,7 +433,9 @@ def predicate(subexpr: itir.Expr, num_occurences: int):
             if num_occurences > 1:
                 if is_local_view:
                     return True
-                else:
+                # condition is only necessary since typing on lambdas is not preserved during
+                #  the transformation
+                elif not isinstance(subexpr, itir.Lambda):
                     # only extract fields outside of `as_fieldop`
                     # `as_fieldop(...)(field_expr, field_expr)`
                     # -> `(λ(_cs_1) → as_fieldop(...)(_cs_1, _cs_1))(field_expr)`
@@ -451,10 +453,8 @@ def predicate(subexpr: itir.Expr, num_occurences: int):
             return self.generic_visit(node, **kwargs)
 
         # apply remapping
-        result = itir.FunCall(
-            fun=itir.Lambda(params=list(extracted.keys()), expr=new_expr),
-            args=list(extracted.values()),
-        )
+        result = im.let(*extracted.items())(new_expr)
+        itir_type_inference.copy_type(from_=node, to=result, allow_untyped=True)
 
         # if the node id is ignored (because its parent is eliminated), but it occurs
         # multiple times then we want to visit the final result once more.

src/gt4py/next/iterator/transforms/inline_lambdas.py

@@ -14,6 +14,7 @@
 from gt4py.next.iterator.ir_utils.common_pattern_matcher import is_applied_lift
 from gt4py.next.iterator.transforms.remap_symbols import RemapSymbolRefs, RenameSymbols
 from gt4py.next.iterator.transforms.symbol_ref_utils import CountSymbolRefs
+from gt4py.next.iterator.type_system import inference as itir_inference
 
 
 # TODO(tehrengruber): Reduce complexity of the function by removing the different options here
@@ -98,7 +99,7 @@ def new_name(name):
         new_expr.location = node.location
         return new_expr
     else:
-        return ir.FunCall(
+        new_expr = ir.FunCall(
             fun=ir.Lambda(
                 params=[
                     param
@@ -110,6 +111,8 @@ def new_name(name):
             args=[arg for arg, eligible in zip(node.args, eligible_params) if not eligible],
             location=node.location,
         )
+        itir_inference.copy_type(from_=node, to=new_expr, allow_untyped=True)
+        return new_expr
 
 
 @dataclasses.dataclass

src/gt4py/next/iterator/type_system/inference.py

@@ -95,14 +95,14 @@ def _set_node_type(node: itir.Node, type_: ts.TypeSpec) -> None:
     node.type = type_
 
 
-def copy_type(from_: itir.Node, to: itir.Node) -> None:
+def copy_type(from_: itir.Node, to: itir.Node, allow_untyped=False) -> None:
     """
     Copy type from one node to another.
 
     This function mainly exists for readability reasons.
     """
-    assert isinstance(from_.type, ts.TypeSpec)
-    _set_node_type(to, from_.type)
+    assert allow_untyped is not None or isinstance(from_.type, ts.TypeSpec)
+    _set_node_type(to, from_.type)  # type: ignore[arg-type]
 
 
 def on_inferred(callback: Callable, *args: Union[ts.TypeSpec, ObservableTypeSynthesizer]) -> None: