Repository: GridTools/gt4py
Pull request #1259 (2023-05-30 10:12:55)

feature[next]: Save type to annex in ITIR type inference
This PR allows saving the results of type inference in the ITIR nodes. This is particularly useful for passes that modify the tree, but need type information during the process (in my case this is for the temporary extraction pass).

```python
testee = im.lambda_("a")(im.plus("a", im.literal("1", "float32")))
ti.infer(testee, save_to_annex=True)
param_type = testee.params[0].annex.type
assert isinstance(param_type, ti.Val) and param_type.dtype.name == "float32"
```


Patch:
diff --git a/AUTHORS.md b/AUTHORS.md
index 7bb08ae7..89aafb99 100644
--- a/AUTHORS.md
+++ b/AUTHORS.md
@@ -8,6 +8,7 @@
 - Davis, Eddie. Allen Institute for AI
 - Deconinck, Florian. SSAI/NASA-GSFC
 - Ehrengruber, Till. ETH Zurich - CSCS
+- Elbert, Oliver D. NOAA-GFDL
 - Farabullini, Nicoletta. ETH Zurich - C2SM
 - George, Rhea. Allen Institute for AI
 - González Paredes, Enrique. ETH Zurich - CSCS
diff --git a/src/gt4py/cartesian/frontend/defir_to_gtir.py b/src/gt4py/cartesian/frontend/defir_to_gtir.py
index 330fa441..f2ee5449 100644
--- a/src/gt4py/cartesian/frontend/defir_to_gtir.py
+++ b/src/gt4py/cartesian/frontend/defir_to_gtir.py
@@ -331,6 +331,7 @@ class DefIRToGTIR(IRNodeVisitor):
         NativeFunction.SQRT: common.NativeFunction.SQRT,
         NativeFunction.EXP: common.NativeFunction.EXP,
         NativeFunction.LOG: common.NativeFunction.LOG,
+        NativeFunction.LOG10: common.NativeFunction.LOG10,
         NativeFunction.GAMMA: common.NativeFunction.GAMMA,
         NativeFunction.CBRT: common.NativeFunction.CBRT,
         NativeFunction.ISFINITE: common.NativeFunction.ISFINITE,
diff --git a/src/gt4py/cartesian/frontend/gtscript_frontend.py b/src/gt4py/cartesian/frontend/gtscript_frontend.py
index cb991e19..1c96773f 100644
--- a/src/gt4py/cartesian/frontend/gtscript_frontend.py
+++ b/src/gt4py/cartesian/frontend/gtscript_frontend.py
@@ -742,6 +742,7 @@ class IRMaker(ast.NodeVisitor):
             "sqrt": nodes.NativeFunction.SQRT,
             "exp": nodes.NativeFunction.EXP,
             "log": nodes.NativeFunction.LOG,
+            "log10": nodes.NativeFunction.LOG10,
             "gamma": nodes.NativeFunction.GAMMA,
             "cbrt": nodes.NativeFunction.CBRT,
             "isfinite": nodes.NativeFunction.ISFINITE,
diff --git a/src/gt4py/cartesian/frontend/nodes.py b/src/gt4py/cartesian/frontend/nodes.py
index f964910b..848f7885 100644
--- a/src/gt4py/cartesian/frontend/nodes.py
+++ b/src/gt4py/cartesian/frontend/nodes.py
@@ -46,7 +46,7 @@ BinaryOperator enumeration (:class:`BinaryOperator`)
 NativeFunction enumeration (:class:`NativeFunction`)
     Native function identifier
     [`ABS`, `MAX`, `MIN, `MOD`, `SIN`, `COS`, `TAN`, `ARCSIN`, `ARCCOS`, `ARCTAN`,
-    `SQRT`, `EXP`, `LOG`, `ISFINITE`, `ISINF`, `ISNAN`, `FLOOR`, `CEIL`, `TRUNC`]
+    `SQRT`, `EXP`, `LOG`, `LOG10`, `ISFINITE`, `ISINF`, `ISNAN`, `FLOOR`, `CEIL`, `TRUNC`]
 
 LevelMarker enumeration (:class:`LevelMarker`)
     Special axis levels
@@ -413,6 +413,7 @@ class NativeFunction(enum.Enum):
     SQRT = enum.auto()
     EXP = enum.auto()
     LOG = enum.auto()
+    LOG10 = enum.auto()
     GAMMA = enum.auto()
     CBRT = enum.auto()
 
@@ -448,6 +449,7 @@ NativeFunction.IR_OP_TO_NUM_ARGS = {
     NativeFunction.SQRT: 1,
     NativeFunction.EXP: 1,
     NativeFunction.LOG: 1,
+    NativeFunction.LOG10: 1,
     NativeFunction.GAMMA: 1,
     NativeFunction.CBRT: 1,
     NativeFunction.ISFINITE: 1,
diff --git a/src/gt4py/cartesian/gtc/common.py b/src/gt4py/cartesian/gtc/common.py
index c2a16418..e1449f11 100644
--- a/src/gt4py/cartesian/gtc/common.py
+++ b/src/gt4py/cartesian/gtc/common.py
@@ -173,6 +173,7 @@ class NativeFunction(eve.StrEnum):
     POW = "pow"
     EXP = "exp"
     LOG = "log"
+    LOG10 = "log10"
     GAMMA = "gamma"
     CBRT = "cbrt"
 
@@ -213,6 +214,7 @@ NativeFunction.IR_OP_TO_NUM_ARGS = {
         NativeFunction.POW: 2,
         NativeFunction.EXP: 1,
         NativeFunction.LOG: 1,
+        NativeFunction.LOG10: 1,
         NativeFunction.GAMMA: 1,
         NativeFunction.CBRT: 1,
         NativeFunction.ISFINITE: 1,
@@ -888,6 +890,7 @@ OP_TO_UFUNC_NAME: Final[
         NativeFunction.POW: "power",
         NativeFunction.EXP: "exp",
         NativeFunction.LOG: "log",
+        NativeFunction.LOG10: "log10",
         NativeFunction.GAMMA: "gamma",
         NativeFunction.CBRT: "cbrt",
         NativeFunction.ISFINITE: "isfinite",
diff --git a/src/gt4py/cartesian/gtc/cuir/cuir_codegen.py b/src/gt4py/cartesian/gtc/cuir/cuir_codegen.py
index 5dd84b7b..142c721a 100644
--- a/src/gt4py/cartesian/gtc/cuir/cuir_codegen.py
+++ b/src/gt4py/cartesian/gtc/cuir/cuir_codegen.py
@@ -169,6 +169,7 @@ class CUIRCodegen(codegen.TemplatedGenerator, eve.VisitorWithSymbolTableTrait):
         NativeFunction.POW: "std::pow",
         NativeFunction.EXP: "std::exp",
         NativeFunction.LOG: "std::log",
+        NativeFunction.LOG10: "std::log10",
         NativeFunction.GAMMA: "std::tgamma",
         NativeFunction.CBRT: "std::cbrt",
         NativeFunction.ISFINITE: "std::isfinite",
diff --git a/src/gt4py/cartesian/gtc/dace/expansion/tasklet_codegen.py b/src/gt4py/cartesian/gtc/dace/expansion/tasklet_codegen.py
index 895e6d01..779fca0c 100644
--- a/src/gt4py/cartesian/gtc/dace/expansion/tasklet_codegen.py
+++ b/src/gt4py/cartesian/gtc/dace/expansion/tasklet_codegen.py
@@ -157,6 +157,7 @@ class TaskletCodegen(eve.codegen.TemplatedGenerator, eve.VisitorWithSymbolTableT
                 common.NativeFunction.POW: "dace.math.pow",
                 common.NativeFunction.EXP: "dace.math.exp",
                 common.NativeFunction.LOG: "dace.math.log",
+                common.NativeFunction.LOG10: "log10",
                 common.NativeFunction.GAMMA: "tgamma",
                 common.NativeFunction.CBRT: "cbrt",
                 common.NativeFunction.ISFINITE: "isfinite",
diff --git a/src/gt4py/cartesian/gtc/gtcpp/gtcpp_codegen.py b/src/gt4py/cartesian/gtc/gtcpp/gtcpp_codegen.py
index 699f15c0..00c3fb16 100644
--- a/src/gt4py/cartesian/gtc/gtcpp/gtcpp_codegen.py
+++ b/src/gt4py/cartesian/gtc/gtcpp/gtcpp_codegen.py
@@ -173,6 +173,7 @@ class GTCppCodegen(codegen.TemplatedGenerator, eve.VisitorWithSymbolTableTrait):
                 NativeFunction.POW: "std::pow",
                 NativeFunction.EXP: "std::exp",
                 NativeFunction.LOG: "std::log",
+                NativeFunction.LOG10: "std::log10",
                 NativeFunction.GAMMA: "std::tgamma",
                 NativeFunction.CBRT: "std::cbrt",
                 NativeFunction.ISFINITE: "std::isfinite",
diff --git a/src/gt4py/cartesian/gtc/passes/oir_optimizations/horizontal_execution_merging.py b/src/gt4py/cartesian/gtc/passes/oir_optimizations/horizontal_execution_merging.py
index d6beed55..9bb5cccc 100644
--- a/src/gt4py/cartesian/gtc/passes/oir_optimizations/horizontal_execution_merging.py
+++ b/src/gt4py/cartesian/gtc/passes/oir_optimizations/horizontal_execution_merging.py
@@ -235,6 +235,7 @@ class OnTheFlyMerging(eve.NodeTranslator, eve.VisitorWithSymbolTableTrait):
                 nf.SQRT,
                 nf.EXP,
                 nf.LOG,
+                nf.LOG10,
                 nf.GAMMA,
                 nf.CBRT,
             }
diff --git a/src/gt4py/cartesian/gtc/ufuncs.py b/src/gt4py/cartesian/gtc/ufuncs.py
index 980abe6b..e3bc5a45 100644
--- a/src/gt4py/cartesian/gtc/ufuncs.py
+++ b/src/gt4py/cartesian/gtc/ufuncs.py
@@ -60,6 +60,7 @@ sqrt: np.ufunc = np.sqrt
 power: np.ufunc = np.power
 exp: np.ufunc = np.exp
 log: np.ufunc = np.log
+log10: np.ufunc = np.log10
 gamma: np.ufunc = gamma_
 cbrt: np.ufunc = np.cbrt
 isfinite: np.ufunc = np.isfinite
diff --git a/src/gt4py/cartesian/gtscript.py b/src/gt4py/cartesian/gtscript.py
index 174e1e23..09e6f0f4 100644
--- a/src/gt4py/cartesian/gtscript.py
+++ b/src/gt4py/cartesian/gtscript.py
@@ -56,6 +56,7 @@ MATH_BUILTINS = {
     "sqrt",
     "exp",
     "log",
+    "log10",
     "gamma",
     "cbrt",
     "isfinite",
@@ -859,6 +860,11 @@ def log(x):
     pass
 
 
+def log10(x):
+    """Return the base-10 logarithm of x."""
+    pass
+
+
 def gamma(x):
     """Return the gamma function of x."""
     pass
diff --git a/src/gt4py/next/ffront/decorator.py b/src/gt4py/next/ffront/decorator.py
index 25d194e1..90b821f8 100644
--- a/src/gt4py/next/ffront/decorator.py
+++ b/src/gt4py/next/ffront/decorator.py
@@ -239,21 +239,10 @@ class Program:
             )
 
     def with_backend(self, backend: ppi.ProgramExecutor) -> Program:
-        return Program(
-            past_node=self.past_node,
-            closure_vars=self.closure_vars,
-            backend=backend,
-            definition=self.definition,  # type: ignore[arg-type]  # mypy wrongly deduces definition as method here
-        )
+        return dataclasses.replace(self, backend=backend)
 
     def with_grid_type(self, grid_type: GridType) -> Program:
-        return Program(
-            past_node=self.past_node,
-            closure_vars=self.closure_vars,
-            backend=self.backend,
-            definition=self.definition,
-            grid_type=grid_type,
-        )
+        return dataclasses.replace(self, grid_type=grid_type)
 
     @functools.cached_property
     def _all_closure_vars(self) -> dict[str, Any]:
@@ -506,22 +495,10 @@ class FieldOperator(GTCallable, Generic[OperatorNodeT]):
         return type_
 
     def with_backend(self, backend: ppi.ProgramExecutor) -> FieldOperator:
-        return FieldOperator(
-            foast_node=self.foast_node,
-            closure_vars=self.closure_vars,
-            definition=self.definition,
-            backend=backend,
-            grid_type=self.grid_type,
-        )
+        return dataclasses.replace(self, backend=backend)
 
-    def with_grid_type(self, grid_type: GridType):
-        return FieldOperator(
-            foast_node=self.foast_node,
-            closure_vars=self.closure_vars,
-            definition=self.definition,
-            backend=self.backend,
-            grid_type=grid_type,
-        )
+    def with_grid_type(self, grid_type: GridType) -> FieldOperator:
+        return dataclasses.replace(self, grid_type=grid_type)
 
     def __gt_itir__(self) -> itir.FunctionDefinition:
         if hasattr(self, "__cached_itir"):
diff --git a/src/gt4py/next/iterator/ir.py b/src/gt4py/next/iterator/ir.py
index 54b06ca8..a28f40e9 100644
--- a/src/gt4py/next/iterator/ir.py
+++ b/src/gt4py/next/iterator/ir.py
@@ -172,24 +172,28 @@ INTEGER_BUILTINS = {"int", "int32", "int64"}
 FLOATING_POINT_BUILTINS = {"float", "float32", "float64"}
 TYPEBUILTINS = {*INTEGER_BUILTINS, *FLOATING_POINT_BUILTINS, "bool"}
 
-BUILTINS = {
+GRAMMAR_BUILTINS = {
     "cartesian_domain",
     "unstructured_domain",
-    "named_range",
+    "make_tuple",
+    "tuple_get",
+    "shift",
     "neighbors",
+    "cast_",
+}
+
+BUILTINS = {
+    *GRAMMAR_BUILTINS,
+    "named_range",
     "list_get",
     "map_",
     "make_const_list",
     "lift",
-    "make_tuple",
-    "tuple_get",
     "reduce",
     "deref",
     "can_deref",
-    "shift",
     "scan",
     "if_",
-    "cast_",
     "floordiv",  # TODO see https://github.com/GridTools/gt4py/issues/1136
     *UNARY_MATH_NUMBER_BUILTINS,
     *UNARY_LOGICAL_BUILTINS,
diff --git a/src/gt4py/next/iterator/transforms/cse.py b/src/gt4py/next/iterator/transforms/cse.py
index a030d61b..e3338676 100644
--- a/src/gt4py/next/iterator/transforms/cse.py
+++ b/src/gt4py/next/iterator/transforms/cse.py
@@ -13,6 +13,9 @@
 # SPDX-License-Identifier: GPL-3.0-or-later
 
 import dataclasses
+import functools
+import operator
+import typing
 
 from gt4py.eve import NodeTranslator, NodeVisitor, SymbolTableTrait, VisitorWithSymbolTableTrait
 from gt4py.eve.utils import UIDGenerator
@@ -24,13 +27,13 @@ from gt4py.next.iterator.transforms.inline_lambdas import inline_lambda
 class _NodeReplacer(NodeTranslator):
     expr_map: dict[int, ir.SymRef]
 
-    def visit_Expr(self, node):
+    def visit_Expr(self, node: ir.Node) -> ir.Node:
         if id(node) in self.expr_map:
             return self.expr_map[id(node)]
         return self.generic_visit(node)
 
-    def visit_FunCall(self, node: ir.FunCall):
-        node = self.visit_Expr(node)
+    def visit_FunCall(self, node: ir.FunCall) -> ir.Node:
+        node = typing.cast(ir.FunCall, self.visit_Expr(node))
         # If we encounter an expression like:
         #  (λ(_cs_1) → (λ(a) → a+a)(_cs_1))(outer_expr)
         # (non-recursively) inline the lambda to obtain:
@@ -49,7 +52,7 @@ class _NodeReplacer(NodeTranslator):
         return node
 
 
-def _is_collectable_expr(node: ir.Node):
+def _is_collectable_expr(node: ir.Node) -> bool:
     if isinstance(node, ir.FunCall):
         # do not collect (and thus deduplicate in CSE) shift(offsets…) calls. Node must still be
         #  visited, to ensure symbol dependencies are recognized correctly.
@@ -66,51 +69,92 @@ def _is_collectable_expr(node: ir.Node):
     return False
 
 
-def _is_if_can_deref(node: ir.Node):
-    # `if_(can_deref(...), ..., ...)`
-    return (
-        isinstance(node, ir.FunCall)
-        and node.fun == ir.SymRef(id="if_")
-        and isinstance(node.args[0], ir.FunCall)
-        and node.args[0].fun == ir.SymRef(id="can_deref")
-    )
-
-
 @dataclasses.dataclass
 class CollectSubexpressions(VisitorWithSymbolTableTrait, NodeVisitor):
-    subexprs: dict[ir.Node, list[tuple[int, set[int]]]] = dataclasses.field(
-        init=False, repr=False, default_factory=dict
-    )
+    @dataclasses.dataclass
+    class State:
+        #: A dictionary mapping a node to a list of node ids which are equal. Additionally, for
+        #: each (actual) node we store a set of all ids of collected child subexpressions.
+        subexprs: dict[ir.Node, list[tuple[int, set[int]]]] = dataclasses.field(
+            default_factory=dict
+        )
+        # TODO(tehrengruber): Revisit if this makes sense or if we can just recompute the collected
+        #  child node ids and get simpler code.
+        #: The ids of all child subexpressions which are collected.
+        collected_child_node_ids: set[int] = dataclasses.field(default_factory=set)
+        #: The ids of all nodes declaring a symbol which are referenced (using a `SymRef`)
+        used_symbol_ids: set[int] = dataclasses.field(default_factory=set)
+
+        def remove_subexprs(self, nodes: typing.Iterable[ir.Node]) -> None:
+            node_ids_to_remove: set[int] = set()
+            for node in nodes:
+                subexpr_data = self.subexprs.pop(node, None)
+                if subexpr_data:
+                    node_ids, _ = zip(*subexpr_data)
+                    node_ids_to_remove |= set(node_ids)
+            for subexpr_data in self.subexprs.values():
+                for _, collected_child_node_ids in subexpr_data:
+                    collected_child_node_ids -= node_ids_to_remove
 
     @classmethod
-    def apply(cls, node: ir.Node):
+    def apply(cls, node: ir.Node) -> dict[ir.Node, list[tuple[int, set[int]]]]:
+        state = cls.State()
         obj = cls()
-        obj.visit(
-            node, used_symbol_ids=set(), collected_child_node_ids=set(), allow_collection=True
-        )
+        obj.visit(node, state=state)
         # return subexpression in pre-order of the tree, i.e. the nodes closer to the root come
         # first, and skip the root node itself
-        return {k: v for k, v in reversed(obj.subexprs.items()) if k is not node}
-
-    def visit(self, node, **kwargs):
-        # TODO(tehrengruber): improve this case as we might miss subexpression that could be eliminated
-        # disable collection (for all child nodes) if node matches `if_(can_deref(...), ..., ...)`
-        if _is_if_can_deref(node):
-            kwargs["allow_collection"] = False
+        return {k: v for k, v in reversed(state.subexprs.items()) if k is not node}
 
+    def visit(self, node: ir.Node, **kwargs) -> None:  # type: ignore[override]  # supertype accepts any node, but we want to be more specific here.
         if not isinstance(node, SymbolTableTrait) and not _is_collectable_expr(node):
             return super().visit(node, **kwargs)
 
-        parents_collected_child_node_ids = kwargs.pop("collected_child_node_ids")
-        parents_used_symbol_ids = kwargs.pop("used_symbol_ids")
-        used_symbol_ids = set[int]()
-        collected_child_node_ids = set[int]()
-        super().visit(
-            node,
-            used_symbol_ids=used_symbol_ids,
-            collected_child_node_ids=collected_child_node_ids,
-            **kwargs,
-        )
+        parent_state = kwargs.pop("state")
+        collected_child_node_ids: set[int] = set()
+        used_symbol_ids: set[int] = set()
+
+        # Special handling of `if_(condition, true_branch, false_branch)` like expressions that
+        # avoids extracting subexpressions unless they are used in at least two of the three
+        # arguments.
+        if isinstance(node, ir.FunCall) and node.fun == ir.SymRef(id="if_"):
+            assert len(node.args) == 3
+            # collect subexpressions for all arguments to the `if_`
+            arg_states = [self.State() for _ in node.args]
+            for arg, state in zip(node.args, arg_states):
+                self.visit(arg, state=state, **kwargs)
+
+            # for each subexpression find in how many of the three arguments they occur
+            subexpr_count: dict[ir.Node, int] = {}
+            for arg_state in arg_states:
+                for subexpr in arg_state.subexprs.keys():
+                    subexpr_count.setdefault(subexpr, 0)
+                    subexpr_count[subexpr] += 1
+
+            # remove all subexpressions that are not eligible for collection
+            eligible_subexprs = {subexpr for subexpr, count in subexpr_count.items() if count >= 2}
+            for arg_state in arg_states:
+                arg_state.remove_subexprs(arg_state.subexprs.keys() - eligible_subexprs)
+
+            # merge the states of the three arguments
+            subexprs: dict[ir.Node, list[tuple[int, set[int]]]] = {}
+            for state in arg_states:
+                for subexpr, data in state.subexprs.items():
+                    subexprs.setdefault(subexpr, []).extend(data)
+            collected_child_node_ids = functools.reduce(
+                operator.or_, (state.collected_child_node_ids for state in arg_states)
+            )
+            used_symbol_ids = functools.reduce(
+                operator.or_, (state.used_symbol_ids for state in arg_states)
+            )
+            # propagate collected subexpressions to parent
+            for subexpr, data in subexprs.items():
+                parent_state.subexprs.setdefault(subexpr, []).extend(data)
+        else:
+            super().visit(
+                node,
+                state=self.State(parent_state.subexprs, collected_child_node_ids, used_symbol_ids),
+                **kwargs,
+            )
 
         if isinstance(node, SymbolTableTrait):
             # remove symbols used in child nodes if they are declared in the current node
@@ -118,22 +162,24 @@ class CollectSubexpressions(VisitorWithSymbolTableTrait, NodeVisitor):
 
         # if no symbols are used that are defined in the root node, i.e. the node given to `apply`,
         # we collect the subexpression
-        if not used_symbol_ids and _is_collectable_expr(node) and kwargs["allow_collection"]:
-            self.subexprs.setdefault(node, []).append((id(node), collected_child_node_ids))
+        if not used_symbol_ids and _is_collectable_expr(node):
+            parent_state.subexprs.setdefault(node, []).append((id(node), collected_child_node_ids))
 
             # propagate to parent that we have collected its child
-            parents_collected_child_node_ids.add(id(node))
+            parent_state.collected_child_node_ids.add(id(node))
 
         # propagate used symbol ids to parent
-        parents_used_symbol_ids.update(used_symbol_ids)
+        parent_state.used_symbol_ids.update(used_symbol_ids)
 
         # propagate to parent which of its children we have collected
         # TODO(tehrengruber): This is expensive for a large tree. Use something like a "ChainSet".
-        parents_collected_child_node_ids.update(collected_child_node_ids)
+        parent_state.collected_child_node_ids.update(collected_child_node_ids)
 
-    def visit_SymRef(self, node: ir.SymRef, *, symtable, used_symbol_ids, **kwargs):
+    def visit_SymRef(
+        self, node: ir.SymRef, *, symtable: dict[str, ir.Node], state: State, **kwargs
+    ) -> None:
         if node.id in symtable:  # root symbol otherwise
-            used_symbol_ids.add(id(symtable[node.id]))
+            state.used_symbol_ids.add(id(symtable[node.id]))
 
 
 @dataclasses.dataclass(frozen=True)
diff --git a/src/gt4py/next/iterator/type_inference.py b/src/gt4py/next/iterator/type_inference.py
index b609ea81..6dc1727b 100644
--- a/src/gt4py/next/iterator/type_inference.py
+++ b/src/gt4py/next/iterator/type_inference.py
@@ -589,14 +589,7 @@ class _TypeInferrer(eve.traits.VisitorWithSymbolTableTrait, eve.NodeTranslator):
         if node.id in ir.BUILTINS:
             if node.id in BUILTIN_TYPES:
                 return freshen(BUILTIN_TYPES[node.id])
-            elif node.id in (
-                "make_tuple",
-                "tuple_get",
-                "shift",
-                "cartesian_domain",
-                "unstructured_domain",
-                "cast_",
-            ):
+            elif node.id in ir.GRAMMAR_BUILTINS:
                 raise TypeError(
                     f"Builtin '{node.id}' is only allowed as applied/called function by the type "
                     f"inference."
@@ -659,6 +652,7 @@ class _TypeInferrer(eve.traits.VisitorWithSymbolTableTrait, eve.NodeTranslator):
             raise TypeError("`tuple_get` requires exactly two arguments.")
         if not isinstance(node.args[0], ir.Literal) or node.args[0].type != "int":
             raise TypeError("The first argument to `tuple_get` must be a literal int.")
+        self.visit(node.args[0], **kwargs)  # visit index so that its type is collected
         idx = int(node.args[0].value)
         tup = self.visit(node.args[1], **kwargs)
         kind = TypeVar.fresh()  # `kind == Iterator()` means splitting an iterator of tuples
@@ -902,15 +896,34 @@ class _TypeInferrer(eve.traits.VisitorWithSymbolTableTrait, eve.NodeTranslator):
         )
 
 
+def _save_types_to_annex(node: ir.Node, types: dict[int, Type]) -> None:
+    for child_node in node.pre_walk_values().if_isinstance(*TYPED_IR_NODES):
+        try:
+            child_node.annex.type = types[id(child_node)]  # type: ignore[attr-defined]
+        except KeyError:
+            if not (
+                isinstance(child_node, ir.SymRef)
+                and child_node.id in ir.GRAMMAR_BUILTINS | ir.TYPEBUILTINS
+            ):
+                raise AssertionError(
+                    f"Expected a type to be inferred for node `{child_node}`, but none was found."
+                )
+
+
 def infer_all(
     node: ir.Node,
+    *,
     offset_provider: Optional[dict[str, Connectivity | Dimension]] = None,
     reindex: bool = True,
+    save_to_annex=False,
 ) -> dict[int, Type]:
     """
     Infer the types of the child expressions of a given iterator IR expression.
 
     The result is a dictionary mapping the (Python) id of child nodes to their type.
+
+    The `save_to_annex` flag should only be used as a last resort when the  return dictionary is
+    not enough.
     """
     # Collect preliminary types of all nodes and constraints on them
     inferrer = _TypeInferrer(offset_provider=offset_provider)
@@ -925,15 +938,24 @@ def infer_all(
     if reindex:
         unified_types = reindex_vars(list(unified_types))
 
-    return {id_: unified_type for id_, unified_type in zip(collected_types.keys(), unified_types)}
+    result = {
+        id_: unified_type
+        for id_, unified_type in zip(collected_types.keys(), unified_types, strict=True)
+    }
+
+    if save_to_annex:
+        _save_types_to_annex(node, result)
+
+    return result
 
 
 def infer(
     expr: ir.Node,
     offset_provider: typing.Optional[dict[str, typing.Any]] = None,
+    save_to_annex: bool = False,
 ) -> Type:
     """Infer the type of the given iterator IR expression."""
-    inferred_types = infer_all(expr, offset_provider)
+    inferred_types = infer_all(expr, offset_provider=offset_provider, save_to_annex=save_to_annex)
     return inferred_types[id(expr)]
 
 


Test patch:
diff --git a/tests/cartesian_tests/integration_tests/multi_feature_tests/stencil_definitions.py b/tests/cartesian_tests/integration_tests/multi_feature_tests/stencil_definitions.py
index 0faaa899..79056c29 100644
--- a/tests/cartesian_tests/integration_tests/multi_feature_tests/stencil_definitions.py
+++ b/tests/cartesian_tests/integration_tests/multi_feature_tests/stencil_definitions.py
@@ -39,6 +39,7 @@ from gt4py.cartesian.gtscript import (
     isinf,
     isnan,
     log,
+    log10,
     mod,
     sin,
     sinh,
@@ -129,7 +130,9 @@ def native_functions(field_a: Field3D, field_b: Field3D):
         tanh_res = tanh(acosh_res)
         atanh_res = atanh(tanh_res)
         sqrt_res = afunc(atanh_res)
-        exp_res = exp(sqrt_res)
+        pow10_res = 10 ** (sqrt_res)
+        log10_res = log10(pow10_res)
+        exp_res = exp(log10_res)
         log_res = log(exp_res)
         gamma_res = gamma(log_res)
         cbrt_res = cbrt(gamma_res)
diff --git a/tests/next_tests/unit_tests/iterator_tests/test_type_inference.py b/tests/next_tests/unit_tests/iterator_tests/test_type_inference.py
index f12e6b37..ea806dbc 100644
--- a/tests/next_tests/unit_tests/iterator_tests/test_type_inference.py
+++ b/tests/next_tests/unit_tests/iterator_tests/test_type_inference.py
@@ -817,7 +817,7 @@ def test_fencil_definition_same_closure_input():
             np.empty((0, 2), dtype=np.int64), Dimension("Edge"), Dimension("Vertex"), 2, False
         )
     }
-    inferred_all: dict[int, ti.Type] = ti.infer_all(testee, offset_provider)
+    inferred_all: dict[int, ti.Type] = ti.infer_all(testee, offset_provider=offset_provider)
 
     # validate locations of fencil params
     fencil_param_types = [inferred_all[id(testee.params[i])] for i in range(3)]
@@ -969,6 +969,13 @@ def test_fencil_definition_with_function_definitions():
     )
 
 
+def test_save_types_to_annex():
+    testee = im.lambda_("a")(im.plus("a", im.literal("1", "float32")))
+    ti.infer(testee, save_to_annex=True)
+    param_type = testee.params[0].annex.type
+    assert isinstance(param_type, ti.Val) and param_type.dtype.name == "float32"
+
+
 def test_pformat():
     vs = [ti.TypeVar(idx=i) for i in range(5)]
     assert ti.pformat(vs[0]) == "T₀"
diff --git a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_cse.py b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_cse.py
index dbbbfd97..2349fb8a 100644
--- a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_cse.py
+++ b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_cse.py
@@ -114,14 +114,15 @@ def test_lambda_redef_same_arg_scope():
     def common_expr():
         return im.lambda_("a")(im.plus("a", im.plus(1, 1)))
 
-    # (λ(f1) → (λ(f2) → f1(2) + f2(2))(λ(a) → a + 1))(λ(a) → a + 1)
+    # (λ(f1) → (λ(f2) → f1(2) + f2(2))(λ(a) → a + (1 + 1)))(λ(a) → a + (1 + 1)) + (1 + 1 + (1 + 1))
     testee = im.plus(
         im.let("f1", common_expr())(
             im.let("f2", common_expr())(im.plus(im.call("f1")(2), im.call("f2")(2)))
         ),
         im.plus(im.plus(1, 1), im.plus(1, 1)),
     )
-    # (λ(_cs_1) → (λ(_cs_2) → _cs_2 + _cs_2)(_cs_1(2)))(λ(a) → a + 1)
+    # (λ(_cs_3) → (λ(_cs_1) → (λ(_cs_4) → _cs_4 + _cs_4 + (_cs_3 + _cs_3))(_cs_1(2)))(
+    #   λ(a) → a + _cs_3))(1 + 1)
     expected = im.let("_cs_3", im.plus(1, 1))(
         im.let("_cs_1", im.lambda_("a")(im.plus("a", "_cs_3")))(
             im.let("_cs_4", im.call("_cs_1")(2))(
@@ -133,22 +134,45 @@ def test_lambda_redef_same_arg_scope():
     assert actual == expected
 
 
-def test_if_can_deref():
-    """
-    Test no subexpression is moved outside expressions of the form `if_(can_deref(...), ..., ...)`
-    """
-    # if can_deref(⟪Iₒ, 1ₒ⟫(it)) then ·⟪Iₒ, 1ₒ⟫(it) else ·⟪Iₒ, 1ₒ⟫(it) + ·⟪Iₒ, 1ₒ⟫(it)
+def test_if_can_deref_no_extraction():
+    # Test that a subexpression only occurring in one branch of an `if_` is not moved outside the
+    # if statement. A case using `can_deref` is used here as it is common.
+
+    # if can_deref(⟪Iₒ, 1ₒ⟫(it)) then ·⟪Iₒ, 1ₒ⟫(it) + ·⟪Iₒ, 1ₒ⟫(it) else 1
     testee = im.if_(
         im.call("can_deref")(im.shift("I", 1)("it")),
-        im.deref(im.shift("I", 1)("it")),
-        # use something more involved where a subexpression can still be eliminated
         im.plus(im.deref(im.shift("I", 1)("it")), im.deref(im.shift("I", 1)("it"))),
+        # use something more involved where a subexpression can still be eliminated
+        im.literal("1", "int"),
     )
-    # if can_deref(⟪Iₒ, 1ₒ⟫(it)) then ·⟪Iₒ, 1ₒ⟫(it) else (λ(_cs_1) → _cs_1 + _cs_1)(·⟪Iₒ, 1ₒ⟫(it))
-    expected = im.if_(
+    # (λ(_cs_1) → if can_deref(_cs_1) then (λ(_cs_2) → _cs_2 + _cs_2)(·_cs_1) else 1)(⟪Iₒ, 1ₒ⟫(it))
+    expected = im.let("_cs_1", im.shift("I", 1)("it"))(
+        im.if_(
+            im.call("can_deref")("_cs_1"),
+            im.let("_cs_2", im.deref("_cs_1"))(im.plus("_cs_2", "_cs_2")),
+            im.literal("1", "int"),
+        )
+    )
+
+    actual = CSE().visit(testee)
+    assert actual == expected
+
+
+def test_if_can_deref_eligible_extraction():
+    # Test that a subexpression only occurring in both branches of an `if_` is moved outside the
+    # if statement. A case using `can_deref` is used here as it is common.
+
+    # if can_deref(⟪Iₒ, 1ₒ⟫(it)) then ·⟪Iₒ, 1ₒ⟫(it) else ·⟪Iₒ, 1ₒ⟫(it) + ·⟪Iₒ, 1ₒ⟫(it)
+    testee = im.if_(
         im.call("can_deref")(im.shift("I", 1)("it")),
         im.deref(im.shift("I", 1)("it")),
-        im.let("_cs_1", im.deref(im.shift("I", 1)("it")))(im.plus("_cs_1", "_cs_1")),
+        im.plus(im.deref(im.shift("I", 1)("it")), im.deref(im.shift("I", 1)("it"))),
+    )
+    # (λ(_cs_3) → (λ(_cs_1) → if can_deref(_cs_3) then _cs_1 else _cs_1 + _cs_1)(·_cs_3))(⟪Iₒ, 1ₒ⟫(it))
+    expected = im.let("_cs_3", im.shift("I", 1)("it"))(
+        im.let("_cs_1", im.deref("_cs_3"))(
+            im.if_(im.call("can_deref")("_cs_3"), "_cs_1", im.plus("_cs_1", "_cs_1"))
+        )
     )
 
     actual = CSE().visit(testee)