Repository: GridTools/gt4py
Pull request #1912 (2025-03-13 10:39:53)

fix[next][dace]: Fix array subset on the branches of if-expressions
The issue was discovered by Philip (@philip-paul-mueller).

The lowering of if-expressions has to ensure that all branches (e.g. if/then states) write the same subset to the result array. This was not the case in baseline when one the branches contained a simple expression such as `{dwdx, dwdy}`. The bug was that the full shape of the global fields was copied into the result array, although the result array was allocated on a sub-domain. This caused an out-of-bound write and a segfault.


Patch:
diff --git a/src/gt4py/next/common.py b/src/gt4py/next/common.py
index ce148dee..d2bc090c 100644
--- a/src/gt4py/next/common.py
+++ b/src/gt4py/next/common.py
@@ -69,7 +69,7 @@ class DimensionKind(StrEnum):
         return self.value
 
 
-_DIM_KIND_ORDER = {DimensionKind.HORIZONTAL: 1, DimensionKind.VERTICAL: 2, DimensionKind.LOCAL: 3}
+_DIM_KIND_ORDER = {DimensionKind.HORIZONTAL: 0, DimensionKind.LOCAL: 1, DimensionKind.VERTICAL: 2}
 
 
 def dimension_to_implicit_offset(dim: str) -> str:
@@ -1159,7 +1159,7 @@ def promote_dims(*dims_list: Sequence[Dimension]) -> list[Dimension]:
     Find an ordering of multiple lists of dimensions.
 
     The resulting list contains all unique dimensions from the input lists,
-    sorted first by dims_kind_order, i.e., `Dimension.kind` (`HORIZONTAL` < `VERTICAL` < `LOCAL`) and then
+    sorted first by dims_kind_order, i.e., `Dimension.kind` (`HORIZONTAL` < `LOCAL` < `VERTICAL`) and then
     lexicographically by `Dimension.value`.
 
     Examples:
@@ -1175,9 +1175,9 @@ def promote_dims(*dims_list: Sequence[Dimension]) -> list[Dimension]:
         Traceback (most recent call last):
         ...
         ValueError: Dimensions 'K[vertical], J[horizontal]' are not ordered correctly, expected 'J[horizontal], K[vertical]'.
-        >>> promote_dims([I, K], [J, E2V]) == [I, J, K, E2V]
+        >>> promote_dims([I, K], [J, E2V]) == [I, J, E2V, K]
         True
-        >>> promote_dims([I, E2C], [K, E2V])
+        >>> promote_dims([I, E2C], [E2V, K])
         Traceback (most recent call last):
         ...
         ValueError: There are more than one dimension with DimensionKind 'LOCAL'.
diff --git a/src/gt4py/next/iterator/ir_utils/ir_makers.py b/src/gt4py/next/iterator/ir_utils/ir_makers.py
index 9d77ca46..aae98335 100644
--- a/src/gt4py/next/iterator/ir_utils/ir_makers.py
+++ b/src/gt4py/next/iterator/ir_utils/ir_makers.py
@@ -437,7 +437,7 @@ def as_fieldop(expr: itir.Expr | str, domain: Optional[itir.Expr] = None) -> Cal
     >>> str(as_fieldop(lambda_("it1", "it2")(plus(deref("it1"), deref("it2"))))("field1", "field2"))
     '(⇑(λ(it1, it2) → ·it1 + ·it2))(field1, field2)'
     """
-    from gt4py.next.iterator.ir_utils import domain_utils
+    from gt4py.next.iterator.ir_utils import common_pattern_matcher as cpm, domain_utils
 
     result = call(
         call("as_fieldop")(
@@ -454,7 +454,9 @@ def as_fieldop(expr: itir.Expr | str, domain: Optional[itir.Expr] = None) -> Cal
 
     def _populate_domain_annex_wrapper(*args, **kwargs):
         node = result(*args, **kwargs)
-        if domain:
+        # note: if the domain is not a direct construction, e.g. because it is only a reference
+        # to a domain defined in a let, don't populate the annex
+        if domain and cpm.is_call_to(domain, ("cartesian_domain", "unstructured_domain")):
             node.annex.domain = domain_utils.SymbolicDomain.from_expr(domain)
         return node
 
diff --git a/src/gt4py/next/iterator/ir_utils/misc.py b/src/gt4py/next/iterator/ir_utils/misc.py
index d604d976..e880e05d 100644
--- a/src/gt4py/next/iterator/ir_utils/misc.py
+++ b/src/gt4py/next/iterator/ir_utils/misc.py
@@ -12,7 +12,7 @@ from collections import ChainMap
 from gt4py import eve
 from gt4py.eve import utils as eve_utils
 from gt4py.next.iterator import ir as itir
-from gt4py.next.iterator.ir_utils import ir_makers as im
+from gt4py.next.iterator.ir_utils import common_pattern_matcher as cpm, ir_makers as im
 
 
 @dataclasses.dataclass(frozen=True)
@@ -73,6 +73,82 @@ def is_equal(a: itir.Expr, b: itir.Expr):
     )
 
 
+def canonicalize_as_fieldop(expr: itir.FunCall) -> itir.FunCall:
+    """
+    Canonicalize applied `as_fieldop`s.
+
+    In case the stencil argument is a `deref` wrap it into a lambda such that we have a unified
+    format to work with (e.g. each parameter has a name without the need to special case).
+    """
+    assert cpm.is_applied_as_fieldop(expr)
+
+    stencil = expr.fun.args[0]  # type: ignore[attr-defined]
+    domain = expr.fun.args[1] if len(expr.fun.args) > 1 else None  # type: ignore[attr-defined]
+    if cpm.is_ref_to(stencil, "deref"):
+        stencil = im.lambda_("arg")(im.deref("arg"))
+        new_expr = im.as_fieldop(stencil, domain)(*expr.args)
+
+        return new_expr
+
+    return expr
+
+
+def _remove_let_alias(let_expr: itir.FunCall):
+    assert cpm.is_let(let_expr)
+    is_aliased_let = True
+    for param, arg in zip(let_expr.fun.params, let_expr.args, strict=True):  # type: ignore[attr-defined]  # ensured by cpm.is_let
+        is_aliased_let &= cpm.is_ref_to(arg, param.id)
+    if is_aliased_let:
+        return let_expr.fun.expr  # type: ignore[attr-defined]  # ensured by cpm.is_let
+    return let_expr
+
+
+def unwrap_scan(stencil: itir.Lambda | itir.FunCall):
+    """
+    If given a scan, extract stencil part of its scan pass and a back-transformation into a scan.
+
+    If a regular stencil is given the stencil is left as-is and the back-transformation is the
+    identity function. This function allows treating a scan stencil like a regular stencil during
+    a transformation avoiding the complexity introduced by the different IR format.
+
+    >>> scan = im.call("scan")(
+    ...     im.lambda_("state", "arg")(im.plus("state", im.deref("arg"))), True, 0.0
+    ... )
+    >>> stencil, back_trafo = unwrap_scan(scan)
+    >>> str(stencil)
+    'λ(arg) → state + ·arg'
+    >>> str(back_trafo(stencil))
+    'scan(λ(state, arg) → state + ·arg, True, 0.0)'
+
+    In case a regular stencil is given it is returned as-is:
+
+    >>> deref_stencil = im.lambda_("it")(im.deref("it"))
+    >>> stencil, back_trafo = unwrap_scan(deref_stencil)
+    >>> assert stencil == deref_stencil
+    """
+    if cpm.is_call_to(stencil, "scan"):
+        scan_pass, direction, init = stencil.args
+        assert isinstance(scan_pass, itir.Lambda)
+        # remove scan pass state to be used by caller
+        state_param = scan_pass.params[0]
+        stencil_like = im.lambda_(*scan_pass.params[1:])(scan_pass.expr)
+
+        def restore_scan(transformed_stencil_like: itir.Lambda):
+            new_scan_pass = im.lambda_(state_param, *transformed_stencil_like.params)(
+                _remove_let_alias(
+                    im.call(transformed_stencil_like)(
+                        *(param.id for param in transformed_stencil_like.params)
+                    )
+                )
+            )
+            return im.call("scan")(new_scan_pass, direction, init)
+
+        return stencil_like, restore_scan
+
+    assert isinstance(stencil, itir.Lambda)
+    return stencil, lambda s: s
+
+
 def with_altered_arg(node: itir.FunCall, arg_idx: int, new_arg: itir.Expr | str) -> itir.FunCall:
     """Given a itir.FunCall return a new call with one of its argument replaced."""
     return im.call(node.fun)(
diff --git a/src/gt4py/next/iterator/transforms/collapse_tuple.py b/src/gt4py/next/iterator/transforms/collapse_tuple.py
index 7a10587d..12028382 100644
--- a/src/gt4py/next/iterator/transforms/collapse_tuple.py
+++ b/src/gt4py/next/iterator/transforms/collapse_tuple.py
@@ -12,36 +12,55 @@ import dataclasses
 import enum
 import functools
 import operator
-from typing import Optional
+import re
+from typing import Literal, Optional
 
 from gt4py import eve
 from gt4py.eve import utils as eve_utils
 from gt4py.next import common
-from gt4py.next.iterator import ir
+from gt4py.next.iterator import ir, ir as itir
 from gt4py.next.iterator.ir_utils import (
     common_pattern_matcher as cpm,
     ir_makers as im,
     misc as ir_misc,
 )
-from gt4py.next.iterator.transforms import fixed_point_transformation
+from gt4py.next.iterator.transforms import fixed_point_transformation, inline_lifts
 from gt4py.next.iterator.transforms.inline_lambdas import InlineLambdas, inline_lambda
-from gt4py.next.iterator.type_system import inference as itir_type_inference
+from gt4py.next.iterator.type_system import (
+    inference as itir_type_inference,
+    type_specifications as it_ts,
+)
 from gt4py.next.type_system import type_info, type_specifications as ts
 
 
+def _with_altered_iterator_element_type(
+    type_: it_ts.IteratorType, new_el_type: ts.DataType
+) -> it_ts.IteratorType:
+    return it_ts.IteratorType(
+        position_dims=type_.position_dims, defined_dims=type_.defined_dims, element_type=new_el_type
+    )
+
+
+def _with_altered_iterator_position_dims(
+    type_: it_ts.IteratorType, new_position_dims: list[common.Dimension] | Literal["unknown"]
+) -> it_ts.IteratorType:
+    return it_ts.IteratorType(
+        position_dims=new_position_dims,
+        defined_dims=type_.defined_dims,
+        element_type=type_.element_type,
+    )
+
+
 def _is_trivial_make_tuple_call(node: ir.Expr):
     """Return if node is a `make_tuple` call with all elements `SymRef`s, `Literal`s or tuples thereof."""
     if not cpm.is_call_to(node, "make_tuple"):
         return False
-    if not all(
-        isinstance(arg, (ir.SymRef, ir.Literal)) or _is_trivial_make_tuple_call(arg)
-        for arg in node.args
-    ):
+    if not all(_is_trivial_or_tuple_thereof_expr(arg) for arg in node.args):
         return False
     return True
 
 
-def _is_trivial_or_tuple_thereof_expr(node: ir.Node) -> bool:
+def _is_trivial_or_tuple_thereof_expr(node: itir.Node) -> bool:
     """
     Return `true` if the expr is a trivial expression (`SymRef` or `Literal`) or tuple thereof.
 
@@ -57,7 +76,7 @@ def _is_trivial_or_tuple_thereof_expr(node: ir.Node) -> bool:
     ... )
     True
     """
-    if isinstance(node, (ir.SymRef, ir.Literal)):
+    if isinstance(node, (itir.SymRef, itir.Literal)):
         return True
     if cpm.is_call_to(node, "make_tuple"):
         return all(_is_trivial_or_tuple_thereof_expr(arg) for arg in node.args)
@@ -74,6 +93,19 @@ def _is_trivial_or_tuple_thereof_expr(node: ir.Node) -> bool:
     return False
 
 
+def _flattened_as_fieldop_param_el_name(param: str, idx: int) -> str:
+    prefix = "__ct_flat_el_"
+
+    # keep the original param name, but skip prefix from previous flattenings
+    if param.startswith(prefix):
+        parent_idx, suffix = re.split(r"_(?!\d)", param[len(prefix) :], maxsplit=1)
+        prefix = f"{prefix}{parent_idx}_"
+    else:
+        suffix = param
+
+    return f"{prefix}{idx}_{suffix}"
+
+
 # TODO(tehrengruber): Conceptually the structure of this pass makes sense: Visit depth first,
 #  transform each node until no transformations apply anymore, whenever a node is to be transformed
 #  go through all available transformation and apply them. However the final result here still
@@ -130,6 +162,10 @@ class CollapseTuple(
         PROPAGATE_NESTED_LET = enum.auto()
         #: `let(a, 1)(a)` -> `1` or `let(a, b)(f(a))` -> `f(a)`
         INLINE_TRIVIAL_LET = enum.auto()
+        #: `as_fieldop(λ(t) → ·t[0]+·t[1])({a, b})` -> `as_fieldop(λ(a, b) → ·a+·b)(a, b)`
+        FLATTEN_AS_FIELDOP_ARGS = enum.auto()
+        #: `let(a, b[1])(a)` -> `b[1]`
+        INLINE_TRIVIAL_TUPLE_LET_VAR = enum.auto()
 
         @classmethod
         def all(self) -> CollapseTuple.Transformation:
@@ -144,7 +180,7 @@ class CollapseTuple(
     @classmethod
     def apply(
         cls,
-        node: ir.Node,
+        node: itir.Node,
         *,
         ignore_tuple_size: bool = False,
         remove_letified_make_tuple_elements: bool = True,
@@ -155,7 +191,7 @@ class CollapseTuple(
         # allow sym references without a symbol declaration, mostly for testing
         allow_undeclared_symbols: bool = False,
         uids: Optional[eve_utils.UIDGenerator] = None,
-    ) -> ir.Node:
+    ) -> itir.Node:
         """
         Simplifies `make_tuple`, `tuple_get` calls.
 
@@ -173,7 +209,7 @@ class CollapseTuple(
         offset_provider_type = offset_provider_type or {}
         uids = uids or eve_utils.UIDGenerator()
 
-        if isinstance(node, ir.Program):
+        if isinstance(node, itir.Program):
             within_stencil = False
         assert within_stencil in [
             True,
@@ -212,18 +248,18 @@ class CollapseTuple(
         return super().visit(node, **kwargs)
 
     def transform_collapse_make_tuple_tuple_get(
-        self, node: ir.FunCall, **kwargs
-    ) -> Optional[ir.Node]:
+        self, node: itir.FunCall, **kwargs
+    ) -> Optional[itir.Node]:
         if cpm.is_call_to(node, "make_tuple") and all(
             cpm.is_call_to(arg, "tuple_get") for arg in node.args
         ):
             # `make_tuple(tuple_get(0, t), tuple_get(1, t), ..., tuple_get(N-1,t))` -> `t`
-            assert isinstance(node.args[0], ir.FunCall)
+            assert isinstance(node.args[0], itir.FunCall)
             first_expr = node.args[0].args[1]
 
             for i, v in enumerate(node.args):
-                assert isinstance(v, ir.FunCall)
-                assert isinstance(v.args[0], ir.Literal)
+                assert isinstance(v, itir.FunCall)
+                assert isinstance(v.args[0], itir.Literal)
                 if not (int(v.args[0].value) == i and ir_misc.is_equal(v.args[1], first_expr)):
                     # tuple argument differs, just continue with the rest of the tree
                     return None
@@ -240,11 +276,11 @@ class CollapseTuple(
         return None
 
     def transform_collapse_tuple_get_make_tuple(
-        self, node: ir.FunCall, **kwargs
-    ) -> Optional[ir.Node]:
+        self, node: itir.FunCall, **kwargs
+    ) -> Optional[itir.Node]:
         if (
             cpm.is_call_to(node, "tuple_get")
-            and isinstance(node.args[0], ir.Literal)
+            and isinstance(node.args[0], itir.Literal)
             and cpm.is_call_to(node.args[1], "make_tuple")
         ):
             # `tuple_get(i, make_tuple(e_0, e_1, ..., e_i, ..., e_N))` -> `e_i`
@@ -257,8 +293,8 @@ class CollapseTuple(
             return node.args[1].args[idx]
         return None
 
-    def transform_propagate_tuple_get(self, node: ir.FunCall, **kwargs) -> Optional[ir.Node]:
-        if cpm.is_call_to(node, "tuple_get") and isinstance(node.args[0], ir.Literal):
+    def transform_propagate_tuple_get(self, node: itir.FunCall, **kwargs) -> Optional[itir.Node]:
+        if cpm.is_call_to(node, "tuple_get") and isinstance(node.args[0], itir.Literal):
             # TODO(tehrengruber): extend to general symbols as long as the tail call in the let
             #   does not capture
             # `tuple_get(i, let(...)(make_tuple()))` -> `let(...)(tuple_get(i, make_tuple()))`
@@ -281,12 +317,14 @@ class CollapseTuple(
                 )
         return None
 
-    def transform_letify_make_tuple_elements(self, node: ir.Node, **kwargs) -> Optional[ir.Node]:
+    def transform_letify_make_tuple_elements(
+        self, node: itir.Node, **kwargs
+    ) -> Optional[itir.Node]:
         if cpm.is_call_to(node, "make_tuple"):
             # `make_tuple(expr1, expr1)`
             # -> `let((_tuple_el_1, expr1), (_tuple_el_2, expr2))(make_tuple(_tuple_el_1, _tuple_el_2))`
-            bound_vars: dict[ir.Sym, ir.Expr] = {}
-            new_args: list[ir.Expr] = []
+            bound_vars: dict[itir.Sym, itir.Expr] = {}
+            new_args: list[itir.Expr] = []
             for arg in node.args:
                 if cpm.is_call_to(node, "make_tuple") and not _is_trivial_make_tuple_call(node):
                     el_name = self.uids.sequential_id(prefix="__ct_el")
@@ -301,7 +339,7 @@ class CollapseTuple(
                 )
         return None
 
-    def transform_inline_trivial_make_tuple(self, node: ir.Node, **kwargs) -> Optional[ir.Node]:
+    def transform_inline_trivial_make_tuple(self, node: itir.Node, **kwargs) -> Optional[itir.Node]:
         if cpm.is_let(node):
             # `let(tup, make_tuple(trivial_expr1, trivial_expr2))(foo(tup))`
             #  -> `foo(make_tuple(trivial_expr1, trivial_expr2))`
@@ -310,13 +348,15 @@ class CollapseTuple(
                 return self.visit(inline_lambda(node, eligible_params=eligible_params), **kwargs)
         return None
 
-    def transform_propagate_to_if_on_tuples(self, node: ir.FunCall, **kwargs) -> Optional[ir.Node]:
+    def transform_propagate_to_if_on_tuples(
+        self, node: itir.FunCall, **kwargs
+    ) -> Optional[itir.Node]:
         if kwargs["within_stencil"]:
             # TODO(tehrengruber): This significantly increases the size of the tree. Skip transformation
             #  in local-view for now. Revisit.
             return None
 
-        if isinstance(node, ir.FunCall) and not cpm.is_call_to(node, "if_"):
+        if isinstance(node, itir.FunCall) and not cpm.is_call_to(node, "if_"):
             # TODO(tehrengruber): Only inline if type of branch value is a tuple.
             # Examples:
             # `(if cond then {1, 2} else {3, 4})[0]` -> `if cond then {1, 2}[0] else {3, 4}[0]`
@@ -335,8 +375,8 @@ class CollapseTuple(
         return None
 
     def transform_propagate_to_if_on_tuples_cps(
-        self, node: ir.FunCall, **kwargs
-    ) -> Optional[ir.Node]:
+        self, node: itir.FunCall, **kwargs
+    ) -> Optional[itir.Node]:
         # The basic idea of this transformation is to remove tuples across if-stmts by rewriting
         # the expression in continuation passing style, e.g. something like a tuple reordering
         # ```
@@ -358,7 +398,7 @@ class CollapseTuple(
         # `if True then {2, 1} else {4, 3}`. The examples in the comments below all refer to this
         # tuple reordering example here.
 
-        if not isinstance(node, ir.FunCall) or cpm.is_call_to(node, "if_"):
+        if not isinstance(node, itir.FunCall) or cpm.is_call_to(node, "if_"):
             return None
 
         # The first argument that is eligible also transforms all remaining args (They will be
@@ -430,7 +470,7 @@ class CollapseTuple(
 
         return None
 
-    def transform_propagate_nested_let(self, node: ir.FunCall, **kwargs) -> Optional[ir.Node]:
+    def transform_propagate_nested_let(self, node: itir.FunCall, **kwargs) -> Optional[itir.Node]:
         if cpm.is_let(node):
             # `let((a, let(b, 1)(a_val)))(a)`-> `let(b, 1)(let(a, a_val)(a))`
             outer_vars = {}
@@ -456,14 +496,95 @@ class CollapseTuple(
                 )
         return None
 
-    def transform_inline_trivial_let(self, node: ir.FunCall, **kwargs) -> Optional[ir.Node]:
+    def transform_inline_trivial_let(self, node: itir.FunCall, **kwargs) -> Optional[itir.Node]:
         if cpm.is_let(node):
-            if isinstance(node.fun.expr, ir.SymRef):  # type: ignore[attr-defined]  # ensured by is_let
+            if isinstance(node.fun.expr, itir.SymRef):  # type: ignore[attr-defined]  # ensured by is_let
                 # `let(a, 1)(a)` -> `1`
                 for arg_sym, arg in zip(node.fun.params, node.args):  # type: ignore[attr-defined]  # ensured by is_let
-                    if isinstance(node.fun.expr, ir.SymRef) and node.fun.expr.id == arg_sym.id:  # type: ignore[attr-defined]  # ensured by is_let
+                    if isinstance(node.fun.expr, itir.SymRef) and node.fun.expr.id == arg_sym.id:  # type: ignore[attr-defined]  # ensured by is_let
                         return arg
-            if any(trivial_args := [isinstance(arg, (ir.SymRef, ir.Literal)) for arg in node.args]):
+            if any(
+                trivial_args := [isinstance(arg, (itir.SymRef, itir.Literal)) for arg in node.args]
+            ):
                 return inline_lambda(node, eligible_params=trivial_args)
 
         return None
+
+    def transform_inline_trivial_tuple_let_var(self, node: ir.Node, **kwargs) -> Optional[ir.Node]:
+        if cpm.is_let(node):
+            if any(trivial_args := [_is_trivial_or_tuple_thereof_expr(arg) for arg in node.args]):
+                return inline_lambda(node, eligible_params=trivial_args)
+        return None
+
+    # TODO(tehrengruber): This is a transformation that should be executed before visiting the children. Then
+    #  revisiting the body would not be needed.
+    def transform_flatten_as_fieldop_args(
+        self, node: itir.FunCall, **kwargs
+    ) -> Optional[itir.Node]:
+        # `as_fieldop(λ(t) → ·t[0]+·t[1])({a, b})` -> `as_fieldop(λ(a, b) → ·a+·b)(a, b)`
+        if not cpm.is_applied_as_fieldop(node):
+            return None
+
+        for arg in node.args:
+            itir_type_inference.reinfer(arg)
+
+        if not any(isinstance(arg.type, ts.TupleType) for arg in node.args):
+            return None
+
+        node = ir_misc.canonicalize_as_fieldop(node)
+        stencil, restore_scan = ir_misc.unwrap_scan(
+            node.fun.args[0]  # type: ignore[attr-defined] # ensured by cpm.is_applied_as_fieldop
+        )
+
+        new_body = stencil.expr
+        domain = node.fun.args[1] if len(node.fun.args) > 1 else None  # type: ignore[attr-defined] # ensured by cpm.is_applied_as_fieldop
+        remapped_args: dict[
+            itir.Sym, itir.Expr
+        ] = {}  # contains the arguments that are remapped, e.g. `{a, b}`
+        new_params: list[itir.Sym] = []
+        new_args: list[itir.Expr] = []
+        for param, arg in zip(stencil.params, node.args, strict=True):
+            if isinstance(arg.type, ts.TupleType):
+                ref_to_remapped_arg = im.ref(f"__ct_flat_remapped_{len(remapped_args)}", arg.type)
+                remapped_args[im.sym(ref_to_remapped_arg.id, arg.type)] = arg
+                new_params_inner, lift_params = [], []
+                for i, type_ in enumerate(param.type.element_type.types):
+                    new_param = im.sym(
+                        _flattened_as_fieldop_param_el_name(param.id, i),
+                        _with_altered_iterator_element_type(param.type, type_),
+                    )
+                    lift_params.append(
+                        im.sym(
+                            new_param.id,
+                            _with_altered_iterator_position_dims(new_param.type, "unknown"),  # type: ignore[arg-type]  # always in IteratorType
+                        )
+                    )
+                    new_params_inner.append(new_param)
+                    new_args.append(im.tuple_get(i, ref_to_remapped_arg))
+
+                # an iterator that substitutes the original (tuple) iterator, e.g. `t`. Built
+                # from the new parameters which are the elements of `t`.
+                param_substitute = im.lift(
+                    im.lambda_(*lift_params)(
+                        im.make_tuple(*[im.deref(im.ref(p.id, p.type)) for p in lift_params])
+                    )
+                )(*[im.ref(p.id, p.type) for p in new_params_inner])
+
+                new_body = im.let(param.id, param_substitute)(new_body)
+                # note: the lift is trivial so inlining it is not an issue with respect to tree size
+                new_body = inline_lambda(new_body, force_inline_lift_args=True)
+
+                new_params.extend(new_params_inner)
+            else:
+                new_params.append(param)
+                new_args.append(arg)
+
+        # remove lifts again
+        new_body = inline_lifts.InlineLifts(
+            flags=inline_lifts.InlineLifts.Flag.INLINE_DEREF_LIFT
+            | inline_lifts.InlineLifts.Flag.PROPAGATE_SHIFT
+        ).visit(new_body)
+        new_body = self.visit(new_body, **kwargs)
+        new_stencil = restore_scan(im.lambda_(*new_params)(new_body))
+
+        return im.let(*remapped_args.items())(im.as_fieldop(new_stencil, domain)(*new_args))
diff --git a/src/gt4py/next/iterator/transforms/fuse_as_fieldop.py b/src/gt4py/next/iterator/transforms/fuse_as_fieldop.py
index 81633dfb..b7463691 100644
--- a/src/gt4py/next/iterator/transforms/fuse_as_fieldop.py
+++ b/src/gt4py/next/iterator/transforms/fuse_as_fieldop.py
@@ -21,6 +21,7 @@ from gt4py.next.iterator.ir_utils import (
     common_pattern_matcher as cpm,
     domain_utils,
     ir_makers as im,
+    misc as ir_misc,
 )
 from gt4py.next.iterator.transforms import (
     fixed_point_transformation,
@@ -46,26 +47,6 @@ def _merge_arguments(
     return new_args
 
 
-def _canonicalize_as_fieldop(expr: itir.FunCall) -> itir.FunCall:
-    """
-    Canonicalize applied `as_fieldop`s.
-
-    In case the stencil argument is a `deref` wrap it into a lambda such that we have a unified
-    format to work with (e.g. each parameter has a name without the need to special case).
-    """
-    assert cpm.is_applied_as_fieldop(expr)
-
-    stencil = expr.fun.args[0]  # type: ignore[attr-defined]
-    domain = expr.fun.args[1] if len(expr.fun.args) > 1 else None  # type: ignore[attr-defined]
-    if cpm.is_ref_to(stencil, "deref"):
-        stencil = im.lambda_("arg")(im.deref("arg"))
-        new_expr = im.as_fieldop(stencil, domain)(*expr.args)
-
-        return new_expr
-
-    return expr
-
-
 def _is_tuple_expr_of_literals(expr: itir.Expr):
     if cpm.is_call_to(expr, "make_tuple"):
         return all(_is_tuple_expr_of_literals(arg) for arg in expr.args)
@@ -78,7 +59,7 @@ def _inline_as_fieldop_arg(
     arg: itir.Expr, *, uids: eve_utils.UIDGenerator
 ) -> tuple[itir.Expr, dict[str, itir.Expr]]:
     assert cpm.is_applied_as_fieldop(arg)
-    arg = _canonicalize_as_fieldop(arg)
+    arg = ir_misc.canonicalize_as_fieldop(arg)
 
     stencil, *_ = arg.fun.args  # type: ignore[attr-defined]  # ensured by `is_applied_as_fieldop`
     inner_args: list[itir.Expr] = arg.args
@@ -114,50 +95,6 @@ def _inline_as_fieldop_arg(
     ), extracted_args
 
 
-def _unwrap_scan(stencil: itir.Lambda | itir.FunCall):
-    """
-    If given a scan, extract stencil part of its scan pass and a back-transformation into a scan.
-
-    If a regular stencil is given the stencil is left as-is and the back-transformation is the
-    identity function. This function allows treating a scan stencil like a regular stencil during
-    a transformation avoiding the complexity introduced by the different IR format.
-
-    >>> scan = im.call("scan")(
-    ...     im.lambda_("state", "arg")(im.plus("state", im.deref("arg"))), True, 0.0
-    ... )
-    >>> stencil, back_trafo = _unwrap_scan(scan)
-    >>> str(stencil)
-    'λ(arg) → state + ·arg'
-    >>> str(back_trafo(stencil))
-    'scan(λ(state, arg) → (λ(arg) → state + ·arg)(arg), True, 0.0)'
-
-    In case a regular stencil is given it is returned as-is:
-
-    >>> deref_stencil = im.lambda_("it")(im.deref("it"))
-    >>> stencil, back_trafo = _unwrap_scan(deref_stencil)
-    >>> assert stencil == deref_stencil
-    """
-    if cpm.is_call_to(stencil, "scan"):
-        scan_pass, direction, init = stencil.args
-        assert isinstance(scan_pass, itir.Lambda)
-        # remove scan pass state to be used by caller
-        state_param = scan_pass.params[0]
-        stencil_like = im.lambda_(*scan_pass.params[1:])(scan_pass.expr)
-
-        def restore_scan(transformed_stencil_like: itir.Lambda):
-            new_scan_pass = im.lambda_(state_param, *transformed_stencil_like.params)(
-                im.call(transformed_stencil_like)(
-                    *(param.id for param in transformed_stencil_like.params)
-                )
-            )
-            return im.call("scan")(new_scan_pass, direction, init)
-
-        return stencil_like, restore_scan
-
-    assert isinstance(stencil, itir.Lambda)
-    return stencil, lambda s: s
-
-
 def fuse_as_fieldop(
     expr: itir.Expr, eligible_args: list[bool], *, uids: eve_utils.UIDGenerator
 ) -> itir.Expr:
@@ -165,7 +102,7 @@ def fuse_as_fieldop(
 
     stencil: itir.Lambda = expr.fun.args[0]  # type: ignore[attr-defined]  # ensured by is_applied_as_fieldop
     assert isinstance(expr.fun.args[0], itir.Lambda) or cpm.is_call_to(stencil, "scan")  # type: ignore[attr-defined]  # ensured by is_applied_as_fieldop
-    stencil, restore_scan = _unwrap_scan(stencil)
+    stencil, restore_scan = ir_misc.unwrap_scan(stencil)
 
     domain = expr.fun.args[1] if len(expr.fun.args) > 1 else None  # type: ignore[attr-defined]  # ensured by is_applied_as_fieldop
 
@@ -411,7 +348,7 @@ class FuseAsFieldOp(
 
     def transform_fuse_as_fieldop(self, node: itir.Node, **kwargs):
         if cpm.is_applied_as_fieldop(node):
-            node = _canonicalize_as_fieldop(node)
+            node = ir_misc.canonicalize_as_fieldop(node)
             stencil = node.fun.args[0]  # type: ignore[attr-defined]  # ensure cpm.is_applied_as_fieldop
             assert isinstance(stencil, itir.Lambda) or cpm.is_call_to(stencil, "scan")
             args: list[itir.Expr] = node.args
diff --git a/src/gt4py/next/program_processors/runners/dace/gtir_builtin_translators.py b/src/gt4py/next/program_processors/runners/dace/gtir_builtin_translators.py
index 6b2a32c0..b0f1cc94 100644
--- a/src/gt4py/next/program_processors/runners/dace/gtir_builtin_translators.py
+++ b/src/gt4py/next/program_processors/runners/dace/gtir_builtin_translators.py
@@ -576,6 +576,105 @@ def translate_as_fieldop(
     )
 
 
+def _construct_if_branch_output(
+    sdfg: dace.SDFG,
+    state: dace.SDFGState,
+    sdfg_builder: gtir_sdfg.SDFGBuilder,
+    domain: gtir.Expr,
+    sym: gtir.Sym,
+    true_br: FieldopData,
+    false_br: FieldopData,
+) -> FieldopData:
+    """
+    Helper function called by `translate_if()` to allocate a temporary field to store
+    the result of an if expression.
+    """
+    assert true_br.gt_type == false_br.gt_type
+    out_type = true_br.gt_type
+
+    if isinstance(sym.type, ts.ScalarType):
+        assert sym.type == out_type
+        dtype = gtx_dace_utils.as_dace_type(sym.type)
+        out, _ = sdfg_builder.add_temp_scalar(sdfg, dtype)
+        out_node = state.add_access(out)
+        return FieldopData(out_node, sym.type, origin=())
+
+    assert isinstance(out_type, ts.FieldType)
+    assert isinstance(sym.type, ts.FieldType)
+    dims, origin, shape = get_field_layout(extract_domain(domain))
+    assert dims == out_type.dims
+
+    if isinstance(out_type.dtype, ts.ScalarType):
+        dtype = gtx_dace_utils.as_dace_type(out_type.dtype)
+    else:
+        assert isinstance(out_type.dtype, ts.ListType)
+        assert out_type.dtype.offset_type is not None
+        assert isinstance(out_type.dtype.element_type, ts.ScalarType)
+        dtype = gtx_dace_utils.as_dace_type(out_type.dtype.element_type)
+        offset_provider_type = sdfg_builder.get_offset_provider_type(
+            out_type.dtype.offset_type.value
+        )
+        assert isinstance(offset_provider_type, gtx_common.NeighborConnectivityType)
+        shape = [*shape, offset_provider_type.max_neighbors]
+
+    out, _ = sdfg_builder.add_temp_array(sdfg, shape, dtype)
+    out_node = state.add_access(out)
+
+    return FieldopData(out_node, out_type, tuple(origin))
+
+
+def _write_if_branch_output(
+    sdfg: dace.SDFG,
+    state: dace.SDFGState,
+    src: FieldopData,
+    dst: FieldopData,
+) -> None:
+    """
+    Helper function called by `translate_if()` to write the result of an if-branch,
+    here `src` field, to the output 'dst' field. The data subset is based on the
+    domain of the `dst` field. Therefore, the full shape of `dst` array is written.
+    """
+    if src.gt_type != dst.gt_type:
+        raise ValueError(
+            f"Source and destination type mismatch, '{dst.gt_type}' vs '{src.gt_type}'."
+        )
+    dst_node = state.add_access(dst.dc_node.data)
+    dst_shape = dst_node.desc(sdfg).shape
+
+    if isinstance(src.gt_type, ts.ScalarType):
+        state.add_nedge(
+            src.dc_node,
+            dst_node,
+            dace.Memlet(data=src.dc_node.data, subset="0"),
+        )
+    else:
+        if isinstance(src.gt_type.dtype, ts.ListType):
+            src_origin = [*src.origin, 0]
+            dst_origin = [*dst.origin, 0]
+        else:
+            src_origin = [*src.origin]
+            dst_origin = [*dst.origin]
+
+        data_subset = dace_subsets.Range(
+            (
+                f"{dst_start - src_start}",
+                f"{dst_start - src_start + size - 1}",  # subtract 1 because the range boundaries are included
+                1,
+            )
+            for src_start, dst_start, size in zip(src_origin, dst_origin, dst_shape, strict=True)
+        )
+
+        state.add_nedge(
+            src.dc_node,
+            dst_node,
+            dace.Memlet(
+                data=src.dc_node.data,
+                subset=data_subset,
+                other_subset=dace_subsets.Range.from_array(dst_node.desc(sdfg)),
+            ),
+        )
+
+
 def translate_if(
     node: gtir.Node,
     sdfg: dace.SDFG,
@@ -590,8 +689,8 @@ def translate_if(
     # expect condition as first argument
     if_stmt = gtir_python_codegen.get_source(cond_expr)
 
-    # use current head state to terminate the dataflow, and add a entry state
-    # to connect the true/false branch states as follows:
+    # evaluate the if-condition in a new entry state and use the current head state
+    # to join the true/false branch states as follows:
     #
     #               ------------
     #           === |   cond   | ===
@@ -618,57 +717,67 @@ def translate_if(
     sdfg.add_edge(cond_state, false_state, dace.InterstateEdge(condition=(f"not ({if_stmt})")))
     sdfg.add_edge(false_state, state, dace.InterstateEdge())
 
-    true_br_args = sdfg_builder.visit(
+    true_br_result = sdfg_builder.visit(
         true_expr,
         sdfg=sdfg,
         head_state=true_state,
     )
-    false_br_args = sdfg_builder.visit(
+    false_br_result = sdfg_builder.visit(
         false_expr,
         sdfg=sdfg,
         head_state=false_state,
     )
 
-    def construct_output(inner_data: FieldopData) -> FieldopData:
-        inner_desc = inner_data.dc_node.desc(sdfg)
-        outer, _ = sdfg_builder.add_temp_array_like(sdfg, inner_desc)
-        outer_node = state.add_access(outer)
-
-        return FieldopData(outer_node, inner_data.gt_type, inner_data.origin)
-
-    result_temps = gtx_utils.tree_map(construct_output)(true_br_args)
-
-    fields: Iterable[tuple[FieldopData, FieldopData, FieldopData]] = zip(
-        gtx_utils.flatten_nested_tuple((true_br_args,)),
-        gtx_utils.flatten_nested_tuple((false_br_args,)),
-        gtx_utils.flatten_nested_tuple((result_temps,)),
-        strict=True,
-    )
-
-    for true_br, false_br, temp in fields:
-        if true_br.gt_type != false_br.gt_type:
-            raise ValueError(
-                f"Different type of result fields on if-branches '{true_br.gt_type}' vs '{false_br.gt_type}'."
+    if isinstance(node.type, ts.TupleType):
+        symbol_tree = gtir_sdfg_utils.make_symbol_tree("x", node.type)
+        if isinstance(node.annex.domain, tuple):
+            domain_tree = node.annex.domain
+        else:
+            # TODO(edopao): this is a workaround for some IR nodes where the inferred
+            #   domain on a tuple of fields is not a tuple, see `test_execution.py::test_ternary_operator_tuple()`
+            domain_tree = gtx_utils.tree_map(lambda _: node.annex.domain)(symbol_tree)
+        node_output = gtx_utils.tree_map(
+            lambda sym,
+            domain,
+            true_br,
+            false_br,
+            sdfg=sdfg,
+            state=state,
+            sdfg_builder=sdfg_builder: _construct_if_branch_output(
+                sdfg,
+                state,
+                sdfg_builder,
+                domain,
+                sym,
+                true_br,
+                false_br,
             )
-        true_br_node = true_br.dc_node
-        false_br_node = false_br.dc_node
-
-        temp_name = temp.dc_node.data
-        true_br_output_node = true_state.add_access(temp_name)
-        true_state.add_nedge(
-            true_br_node,
-            true_br_output_node,
-            sdfg.make_array_memlet(temp_name),
+        )(
+            symbol_tree,
+            domain_tree,
+            true_br_result,
+            false_br_result,
         )
-
-        false_br_output_node = false_state.add_access(temp_name)
-        false_state.add_nedge(
-            false_br_node,
-            false_br_output_node,
-            sdfg.make_array_memlet(temp_name),
+        gtx_utils.tree_map(
+            lambda src, dst, state=true_state: _write_if_branch_output(sdfg, state, src, dst)
+        )(true_br_result, node_output)
+        gtx_utils.tree_map(
+            lambda src, dst, state=false_state: _write_if_branch_output(sdfg, state, src, dst)
+        )(false_br_result, node_output)
+    else:
+        node_output = _construct_if_branch_output(
+            sdfg,
+            state,
+            sdfg_builder,
+            node.annex.domain,
+            im.sym("x", node.type),
+            true_br_result,
+            false_br_result,
         )
+        _write_if_branch_output(sdfg, true_state, true_br_result, node_output)
+        _write_if_branch_output(sdfg, false_state, false_br_result, node_output)
 
-    return result_temps
+    return node_output
 
 
 def translate_index(
@@ -739,7 +848,7 @@ def _get_data_nodes(
             )
         else:
             data_node = state.add_access(data_name)
-        return sdfg_builder.make_field(data_node, data_type)
+        return FieldopData(data_node, data_type, origin=())
 
     elif isinstance(data_type, ts.TupleType):
         symbol_tree = gtir_sdfg_utils.make_symbol_tree(data_name, data_type)
diff --git a/src/gt4py/next/program_processors/runners/dace/gtir_sdfg.py b/src/gt4py/next/program_processors/runners/dace/gtir_sdfg.py
index a4c01948..aefe9f32 100644
--- a/src/gt4py/next/program_processors/runners/dace/gtir_sdfg.py
+++ b/src/gt4py/next/program_processors/runners/dace/gtir_sdfg.py
@@ -122,7 +122,7 @@ class SDFGBuilder(DataflowBuilder, Protocol):
     def make_field(
         self,
         data_node: dace.nodes.AccessNode,
-        data_type: ts.FieldType | ts.ScalarType,
+        data_type: ts.FieldType,
     ) -> gtir_builtin_translators.FieldopData:
         """Retrieve the field data descriptor including the domain offset information."""
         ...
@@ -248,7 +248,7 @@ class GTIRToSDFG(eve.NodeVisitor, SDFGBuilder):
     def make_field(
         self,
         data_node: dace.nodes.AccessNode,
-        data_type: ts.FieldType | ts.ScalarType,
+        data_type: ts.FieldType,
     ) -> gtir_builtin_translators.FieldopData:
         """
         Helper method to build the field data type associated with a data access node.
@@ -271,8 +271,6 @@ class GTIRToSDFG(eve.NodeVisitor, SDFGBuilder):
         Returns:
             The descriptor associated with the SDFG data storage, filled with field origin.
         """
-        if isinstance(data_type, ts.ScalarType):
-            return gtir_builtin_translators.FieldopData(data_node, data_type, origin=())
         local_dims = [dim for dim in data_type.dims if dim.kind == gtx_common.DimensionKind.LOCAL]
         if len(local_dims) == 0:
             # do nothing: the field domain consists of all global dimensions
@@ -495,18 +493,24 @@ class GTIRToSDFG(eve.NodeVisitor, SDFGBuilder):
         assert sink_states[0] == head_state
 
         def make_temps(
-            field: gtir_builtin_translators.FieldopData,
+            src: gtir_builtin_translators.FieldopData,
         ) -> gtir_builtin_translators.FieldopData:
-            desc = sdfg.arrays[field.dc_node.data]
-            if desc.transient or not use_temp:
-                return field
+            src_desc = sdfg.arrays[src.dc_node.data]
+            if src_desc.transient or not use_temp:
+                return src
             else:
-                temp, _ = self.add_temp_array_like(sdfg, desc)
-                temp_node = head_state.add_access(temp)
+                dst, dst_desc = self.add_temp_array_like(sdfg, src_desc)
+                dst_node = head_state.add_access(dst)
                 head_state.add_nedge(
-                    field.dc_node, temp_node, sdfg.make_array_memlet(field.dc_node.data)
+                    src.dc_node,
+                    dst_node,
+                    dace.Memlet(
+                        data=dst,
+                        subset=dace.subsets.Range.from_array(dst_desc),
+                        other_subset=dace.subsets.Range.from_array(src_desc),
+                    ),
                 )
-                return gtir_builtin_translators.FieldopData(temp_node, field.gt_type, field.origin)
+                return gtir_builtin_translators.FieldopData(dst_node, src.gt_type, src.origin)
 
         temp_result = gtx_utils.tree_map(make_temps)(result)
         return list(gtx_utils.flatten_nested_tuple((temp_result,)))


Test patch:
diff --git a/tests/next_tests/unit_tests/ffront_tests/test_type_deduction.py b/tests/next_tests/unit_tests/ffront_tests/test_type_deduction.py
index f9393bd9..d3396a54 100644
--- a/tests/next_tests/unit_tests/ffront_tests/test_type_deduction.py
+++ b/tests/next_tests/unit_tests/ffront_tests/test_type_deduction.py
@@ -163,6 +163,20 @@ def test_premap_nbfield(premap_setup):
     )
 
 
+def test_premap_nbfield_with_vertical(premap_setup):
+    X, Y, Y2XDim, Y2X = premap_setup
+    K = Dimension("K", kind=DimensionKind.VERTICAL)
+
+    def premap_fo(bar: Field[[X, K], int64]) -> Field[[Y, Y2XDim, K], int64]:
+        return bar(Y2X)
+
+    parsed = FieldOperatorParser.apply_to_function(premap_fo)
+
+    assert parsed.body.stmts[0].value.type == ts.FieldType(
+        dims=[Y, Y2XDim, K], dtype=ts.ScalarType(kind=ts.ScalarKind.INT64)
+    )
+
+
 def test_premap_reduce(premap_setup):
     X, Y, Y2XDim, Y2X = premap_setup
 
diff --git a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_collapse_tuple.py b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_collapse_tuple.py
index 916ae4e5..636e6694 100644
--- a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_collapse_tuple.py
+++ b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_collapse_tuple.py
@@ -5,11 +5,14 @@
 #
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
-
+from gt4py.next import common
 from gt4py.next.iterator.ir_utils import ir_makers as im
 from gt4py.next.iterator.transforms.collapse_tuple import CollapseTuple
 from gt4py.next.type_system import type_specifications as ts
-from next_tests.unit_tests.iterator_tests.test_type_inference import int_type
+from gt4py.next.iterator.type_system import type_specifications as it_ts
+
+int_type = ts.ScalarType(kind=ts.ScalarKind.INT32)
+Vertex = common.Dimension(value="Vertex", kind=common.DimensionKind.HORIZONTAL)
 
 
 def test_simple_make_tuple_tuple_get():
@@ -125,8 +128,11 @@ def test_propagate_tuple_get():
 
 
 def test_letify_make_tuple_elements():
-    # anything that is not trivial, i.e. a SymRef, works here
-    el1, el2 = im.let("foo", "foo")("foo"), im.let("bar", "bar")("bar")
+    fun_type = ts.FunctionType(
+        pos_only_args=[], pos_or_kw_args={}, kw_only_args={}, returns=int_type
+    )
+    # anything that is not trivial, works here
+    el1, el2 = im.call(im.ref("foo", fun_type))(), im.call(im.ref("bar", fun_type))()
     testee = im.make_tuple(el1, el2)
     expected = im.let(("__ct_el_1", el1), ("__ct_el_2", el2))(
         im.make_tuple("__ct_el_1", "__ct_el_2")
@@ -311,3 +317,81 @@ def test_if_make_tuple_reorder_cps_external():
         within_stencil=False,
     )
     assert actual == expected
+
+
+def test_flatten_as_fieldop_args():
+    it_type = it_ts.IteratorType(
+        position_dims=[Vertex],
+        defined_dims=[Vertex],
+        element_type=ts.TupleType(types=[int_type, int_type]),
+    )
+    testee = im.as_fieldop(im.lambda_(im.sym("it", it_type))(im.tuple_get(1, im.deref("it"))))(
+        im.make_tuple(1, 2)
+    )
+    expected = im.as_fieldop(
+        im.lambda_("__ct_flat_el_0_it", "__ct_flat_el_1_it")(im.deref("__ct_flat_el_1_it"))
+    )(1, 2)
+    actual = CollapseTuple.apply(
+        testee,
+        enabled_transformations=~CollapseTuple.Transformation.PROPAGATE_TO_IF_ON_TUPLES,
+        allow_undeclared_symbols=True,
+        within_stencil=False,
+    )
+    assert actual == expected
+
+
+def test_flatten_as_fieldop_args_nested():
+    it_type = it_ts.IteratorType(
+        position_dims=[Vertex],
+        defined_dims=[Vertex],
+        element_type=ts.TupleType(
+            types=[
+                int_type,
+                ts.TupleType(types=[int_type, int_type]),
+            ]
+        ),
+    )
+    testee = im.as_fieldop(
+        im.lambda_(im.sym("it", it_type))(im.tuple_get(1, im.tuple_get(1, im.deref("it"))))
+    )(im.make_tuple(1, im.make_tuple(2, 3)))
+    expected = im.as_fieldop(
+        im.lambda_("__ct_flat_el_0_it", "__ct_flat_el_1_0_it", "__ct_flat_el_1_1_it")(
+            im.deref("__ct_flat_el_1_1_it")
+        )
+    )(1, 2, 3)
+    actual = CollapseTuple.apply(
+        testee,
+        enabled_transformations=~CollapseTuple.Transformation.PROPAGATE_TO_IF_ON_TUPLES,
+        allow_undeclared_symbols=True,
+        within_stencil=False,
+    )
+    assert actual == expected
+
+
+def test_flatten_as_fieldop_args_scan():
+    it_type = it_ts.IteratorType(
+        position_dims=[Vertex],
+        defined_dims=[Vertex],
+        element_type=ts.TupleType(types=[int_type, int_type]),
+    )
+    testee = im.as_fieldop(
+        im.scan(
+            im.lambda_("state", im.sym("it", it_type))(im.tuple_get(1, im.deref("it"))), True, 0
+        )
+    )(im.make_tuple(1, 2))
+    expected = im.as_fieldop(
+        im.scan(
+            im.lambda_("state", "__ct_flat_el_0_it", "__ct_flat_el_1_it")(
+                im.deref("__ct_flat_el_1_it")
+            ),
+            True,
+            0,
+        )
+    )(1, 2)
+    actual = CollapseTuple.apply(
+        testee,
+        enabled_transformations=~CollapseTuple.Transformation.PROPAGATE_TO_IF_ON_TUPLES,
+        allow_undeclared_symbols=True,
+        within_stencil=False,
+    )
+    assert actual == expected
diff --git a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_gtir_to_sdfg.py b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_gtir_to_sdfg.py
index 030aa9b1..cefa2e4e 100644
--- a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_gtir_to_sdfg.py
+++ b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_gtir_to_sdfg.py
@@ -98,6 +98,12 @@ def make_mesh_symbols(mesh: MeshDescriptor):
 def build_dace_sdfg(
     ir: gtir.Program, offset_provider_type: gtx_common.OffsetProviderType
 ) -> Callable[..., Any]:
+    # run domain inference in order to add the domain annex information to the IR nodes
+    ir = infer_domain.infer_program(
+        ir,
+        offset_provider=offset_provider_type,
+        symbolic_domain_sizes={IDim.value: "size", Edge.value: "nedges", Vertex.value: "nvertices"},
+    )
     return dace_backend.build_sdfg_from_gtir(
         ir, offset_provider_type, disable_field_origin_on_program_arguments=True
     )
@@ -662,6 +668,7 @@ def test_gtir_cond():
         assert np.allclose(d, (a + b + 1) if s1 > s2 else (a + c + 1))
 
 
+@pytest.mark.xfail(reason="requires function to retrieve the annex tuple domain")
 def test_gtir_cond_with_tuple_return():
     domain = im.domain(gtx_common.GridType.CARTESIAN, ranges={IDim: (0, "size")})
     testee = gtir.Program(
@@ -1844,7 +1851,11 @@ def test_gtir_let_lambda_with_tuple1():
     a = np.random.rand(N)
     b = np.random.rand(N)
 
-    sdfg = build_dace_sdfg(testee, CARTESIAN_OFFSETS)
+    # TODO(edopao): call `build_dace_sdfg` as in all other tests, once this error is fixed
+    #   in domain inference: 'target_domain' cannot be 'NEVER' unless `allow_uninferred=True`
+    sdfg = dace_backend.build_sdfg_from_gtir(
+        testee, CARTESIAN_OFFSETS, disable_field_origin_on_program_arguments=True
+    )
 
     z_fields = (np.zeros_like(a), np.zeros_like(a))
     a_ref = np.concatenate((z_fields[0][:1], a[1 : N - 1], z_fields[0][N - 1 :]))
@@ -2039,8 +2050,6 @@ def test_gtir_index():
 
     v = np.zeros(N, dtype=np.int32)
 
-    # we need to run domain inference in order to add the domain annex information to the index node.
-    testee = infer_domain.infer_program(testee, offset_provider=CARTESIAN_OFFSETS)
     sdfg = build_dace_sdfg(testee, CARTESIAN_OFFSETS)
 
     ref = np.concatenate(
diff --git a/tests/next_tests/unit_tests/test_common.py b/tests/next_tests/unit_tests/test_common.py
index 09ca44aa..c38496eb 100644
--- a/tests/next_tests/unit_tests/test_common.py
+++ b/tests/next_tests/unit_tests/test_common.py
@@ -582,11 +582,11 @@ def dimension_promotion_cases() -> (
             "Dimensions 'KDim[vertical], JDim[horizontal]' are not ordered correctly, expected 'JDim[horizontal], KDim[vertical]'.",
         ),
         (
-            [[JDim, V2E], [IDim, KDim, E2C2V]],
+            [[JDim, V2E], [IDim, E2C2V, KDim]],
             None,
             "There are more than one dimension with DimensionKind 'LOCAL'.",
         ),
-        ([[JDim, V2E], [IDim, KDim]], [IDim, JDim, KDim, V2E], None),
+        ([[JDim, V2E], [IDim, KDim]], [IDim, JDim, V2E, KDim], None),
     ]
     return [
         ([[el for el in arg] for arg in args], [el for el in result] if result else result, msg)