Repository: GridTools/gt4py
Pull request #1079 (2022-12-08 11:36:26)

feature[next]: Scalar if
Adds support for regular if-statements (with boolean condition).
```python
@field_operator(backend=fieldview_backend)
def both_branches_return(inp: Field[[IDim, JDim], float64], condition1: bool, condition2: bool):
    if condition1:
        tmp1 = inp
        if condition2:
            return tmp1 + 1.0
        result = tmp1 + 2.0
    else:
        result = inp + 3.0
    return result
```
~Including extended SSA pass supporting Ifs.~

Superseeds #972

tehrengruber: @egparedes This PR is ready for review (after #1250 is merged). The files `single_static_assign.py` and `test_single_static_assign.py` only appear here until we merge #1250.

Patch:
diff --git a/src/gt4py/eve/traits.py b/src/gt4py/eve/traits.py
index efcf12af..df556c9d 100644
--- a/src/gt4py/eve/traits.py
+++ b/src/gt4py/eve/traits.py
@@ -59,7 +59,14 @@ class SymbolTableTrait:
                         )
                     self.collected_symbols[symbol_name] = node
 
-            if not isinstance(node, SymbolTableTrait):
+            # TODO(egparedes): revisit and generalize mechanism to resolve name collisions.
+            if hasattr(node.annex, "propagated_symbols"):
+                # ensure we have no collisions
+                assert not set(self.collected_symbols.keys()) & set(
+                    node.annex.propagated_symbols.keys()
+                )
+                self.collected_symbols = {**self.collected_symbols, **node.annex.propagated_symbols}
+            elif not isinstance(node, SymbolTableTrait):
                 # Stop recursion if the node opens a new scope (i.e. node with SymbolTableTrait)
                 self.generic_visit(node)
 
diff --git a/src/gt4py/next/ffront/field_operator_ast.py b/src/gt4py/next/ffront/field_operator_ast.py
index 970fe2b4..6b772227 100644
--- a/src/gt4py/next/ffront/field_operator_ast.py
+++ b/src/gt4py/next/ffront/field_operator_ast.py
@@ -16,7 +16,7 @@ from __future__ import annotations
 
 from typing import Any, Generic, TypeVar, Union
 
-from gt4py.eve import Coerced, Node, SourceLocation, SymbolName, SymbolRef
+from gt4py.eve import Coerced, Node, SourceLocation, SymbolName, SymbolRef, datamodels
 from gt4py.eve.traits import SymbolTableTrait
 from gt4py.eve.type_definitions import StrEnum
 from gt4py.next.ffront import dialect_ast_enums, type_specifications as ts_ffront
@@ -179,6 +179,25 @@ class BlockStmt(Stmt, SymbolTableTrait):
     stmts: list[Stmt]
 
 
+class IfStmt(Stmt):
+    condition: Expr
+    true_branch: BlockStmt
+    false_branch: BlockStmt
+
+    @datamodels.root_validator
+    @classmethod
+    def _collect_common_symbols(cls: type[IfStmt], instance: IfStmt) -> None:
+        common_symbol_names = (
+            instance.true_branch.annex.symtable.keys() & instance.false_branch.annex.symtable.keys()
+        )
+        instance.annex.propagated_symbols = {
+            sym_name: Symbol(
+                id=sym_name, type=ts.DeferredType(constraint=None), location=instance.location
+            )
+            for sym_name in common_symbol_names
+        }
+
+
 class FunctionDefinition(LocatedNode, SymbolTableTrait):
     id: Coerced[SymbolName]  # noqa: A003  # shadowing a python builtin
     params: list[DataSymbol]
diff --git a/src/gt4py/next/ffront/foast_introspection.py b/src/gt4py/next/ffront/foast_introspection.py
index 18564438..805df465 100644
--- a/src/gt4py/next/ffront/foast_introspection.py
+++ b/src/gt4py/next/ffront/foast_introspection.py
@@ -19,6 +19,7 @@ from gt4py.next.ffront import field_operator_ast as foast
 
 class StmtReturnKind(enum.IntEnum):
     UNCONDITIONAL_RETURN = 0
+    CONDITIONAL_RETURN = 1
     NO_RETURN = 2
 
 
@@ -33,6 +34,13 @@ def deduce_stmt_return_kind(node: foast.Stmt) -> StmtReturnKind:
         else:
           return 2
 
+    Example with ``StmtReturnKind.CONDITIONAL_RETURN``::
+
+        if cond:
+          return 1
+        else:
+          result = 2
+
     Example with ``StmtReturnKind.NO_RETURN``::
 
         if cond:
@@ -40,7 +48,21 @@ def deduce_stmt_return_kind(node: foast.Stmt) -> StmtReturnKind:
         else:
           result = 2
     """
-    if isinstance(node, foast.Return):
+    if isinstance(node, foast.IfStmt):
+        return_kinds = (
+            deduce_stmt_return_kind(node.true_branch),
+            deduce_stmt_return_kind(node.false_branch),
+        )
+        if all(return_kind is StmtReturnKind.UNCONDITIONAL_RETURN for return_kind in return_kinds):
+            return StmtReturnKind.UNCONDITIONAL_RETURN
+        elif any(
+            return_kind in (StmtReturnKind.UNCONDITIONAL_RETURN, StmtReturnKind.CONDITIONAL_RETURN)
+            for return_kind in return_kinds
+        ):
+            return StmtReturnKind.CONDITIONAL_RETURN
+        assert all(return_kind is StmtReturnKind.NO_RETURN for return_kind in return_kinds)
+        return StmtReturnKind.NO_RETURN
+    elif isinstance(node, foast.Return):
         return StmtReturnKind.UNCONDITIONAL_RETURN
     elif isinstance(node, foast.BlockStmt):
         for stmt in node.stmts:
diff --git a/src/gt4py/next/ffront/foast_passes/type_deduction.py b/src/gt4py/next/ffront/foast_passes/type_deduction.py
index 9f7ac1a1..33ea81ff 100644
--- a/src/gt4py/next/ffront/foast_passes/type_deduction.py
+++ b/src/gt4py/next/ffront/foast_passes/type_deduction.py
@@ -116,18 +116,72 @@ def promote_to_mask_type(
         return input_type
 
 
-def deduce_stmt_return_type(node: foast.BlockStmt) -> ts.TypeSpec:
-    """Deduce type of value returned inside a block statement."""
+def deduce_stmt_return_type(
+    node: foast.BlockStmt, *, requires_unconditional_return=True
+) -> Optional[ts.TypeSpec]:
+    """
+    Deduce type of value returned inside a block statement.
+
+    If `requires_unconditional_return` is true the function additionally ensures that the block
+    statement unconditionally returns and raises an `AssertionError` if not.
+    """
+    conditional_return_type: Optional[ts.TypeSpec] = None
+
     for stmt in node.stmts:
+        is_unconditional_return = False
+        return_type: Optional[ts.TypeSpec]
+
         if isinstance(stmt, foast.Return):
-            return stmt.value.type
+            is_unconditional_return = True
+            return_type = stmt.value.type
+        elif isinstance(stmt, foast.IfStmt):
+            return_types = (
+                deduce_stmt_return_type(stmt.true_branch, requires_unconditional_return=False),
+                deduce_stmt_return_type(stmt.false_branch, requires_unconditional_return=False),
+            )
+            # if both branches return
+            if return_types[0] and return_types[1]:
+                if return_types[0] == return_types[1]:
+                    is_unconditional_return = True
+                else:
+                    raise FieldOperatorTypeDeductionError.from_foast_node(
+                        stmt,
+                        msg=f"If statement contains return statements with inconsistent types:"
+                        f"{return_types[0]} != {return_types[1]}",
+                    )
+            return_type = return_types[0] or return_types[1]
+        elif isinstance(stmt, foast.BlockStmt):
+            # just forward to nested BlockStmt
+            return_type = deduce_stmt_return_type(
+                stmt, requires_unconditional_return=requires_unconditional_return
+            )
+        elif isinstance(stmt, (foast.Assign, foast.TupleTargetAssign)):
+            return_type = None
+        else:
+            raise AssertionError(f"Nodes of type `{type(stmt).__name__}` not supported.")
 
-    # If the node was constructed by the foast parsing we should never get here, but instead
-    # have gotten an error there.
-    raise AssertionError(
-        "Malformed block statement. Expected a return statement in this context, "
-        "but none was found. Please submit a bug report."
-    )
+        if conditional_return_type and return_type and return_type != conditional_return_type:
+            raise FieldOperatorTypeDeductionError.from_foast_node(
+                stmt,
+                msg=f"If statement contains return statements with inconsistent types:"
+                f"{conditional_return_type} != {conditional_return_type}",
+            )
+
+        if is_unconditional_return:  # found a statement that always returns
+            assert return_type
+            return return_type
+        elif return_type:
+            conditional_return_type = return_type
+
+    if requires_unconditional_return:
+        # If the node was constructed by the foast parsing we should never get here, but instead
+        # we should have gotten an error there.
+        raise AssertionError(
+            "Malformed block statement. Expected a return statement in this context, "
+            "but none was found. Please submit a bug report."
+        )
+
+    return None
 
 
 class FieldOperatorTypeDeductionCompletnessValidator(NodeVisitor):
@@ -328,6 +382,48 @@ class FieldOperatorTypeDeduction(traits.VisitorWithSymbolTableTrait, NodeTransla
 
         return foast.TupleTargetAssign(targets=new_targets, value=values, location=node.location)
 
+    def visit_IfStmt(self, node: foast.IfStmt, **kwargs) -> foast.IfStmt:
+        symtable = kwargs["symtable"]
+
+        new_true_branch = self.visit(node.true_branch, **kwargs)
+        new_false_branch = self.visit(node.false_branch, **kwargs)
+        new_node = foast.IfStmt(
+            condition=self.visit(node.condition, **kwargs),
+            true_branch=new_true_branch,
+            false_branch=new_false_branch,
+            location=node.location,
+        )
+
+        if not isinstance(new_node.condition.type, ts.ScalarType):
+            raise FieldOperatorTypeDeductionError.from_foast_node(
+                node,
+                msg="Condition for `if` must be scalar. "
+                f"But got `{new_node.condition.type}` instead.",
+            )
+
+        if new_node.condition.type.kind != ts.ScalarKind.BOOL:
+            raise FieldOperatorTypeDeductionError.from_foast_node(
+                node,
+                msg="Condition for `if` must be of boolean type. "
+                f"But got `{new_node.condition.type}` instead.",
+            )
+
+        for sym in node.annex.propagated_symbols.keys():
+            if (true_type := new_true_branch.annex.symtable[sym].type) != (
+                false_type := new_false_branch.annex.symtable[sym].type
+            ):
+                raise FieldOperatorTypeDeductionError.from_foast_node(
+                    node,
+                    msg=f"Inconsistent types between two branches for variable `{sym}`. "
+                    f"Got types `{true_type}` and `{false_type}.",
+                )
+            # TODO: properly patch symtable (new node?)
+            symtable[sym].type = new_node.annex.propagated_symbols[
+                sym
+            ].type = new_true_branch.annex.symtable[sym].type
+
+        return new_node
+
     def visit_Symbol(
         self,
         node: foast.Symbol,
diff --git a/src/gt4py/next/ffront/foast_to_itir.py b/src/gt4py/next/ffront/foast_to_itir.py
index 2811be1e..05e2d818 100644
--- a/src/gt4py/next/ffront/foast_to_itir.py
+++ b/src/gt4py/next/ffront/foast_to_itir.py
@@ -12,9 +12,11 @@
 #
 # SPDX-License-Identifier: GPL-3.0-or-later
 
-from typing import Any, Callable
+import dataclasses
+from typing import Any, Callable, Optional
 
 from gt4py.eve import NodeTranslator
+from gt4py.eve.utils import UIDGenerator
 from gt4py.next.ffront import (
     dialect_ast_enums,
     fbuiltins,
@@ -22,6 +24,7 @@ from gt4py.next.ffront import (
     type_specifications as ts_ffront,
 )
 from gt4py.next.ffront.fbuiltins import FUN_BUILTIN_NAMES, MATH_BUILTIN_NAMES, TYPE_BUILTIN_NAMES
+from gt4py.next.ffront.foast_introspection import StmtReturnKind, deduce_stmt_return_kind
 from gt4py.next.iterator import ir as itir, ir_makers as im
 from gt4py.next.type_system import type_info, type_specifications as ts
 
@@ -34,6 +37,7 @@ def promote_to_list(
     return lambda x: x
 
 
+@dataclasses.dataclass
 class FieldOperatorLowering(NodeTranslator):
     """
     Lower FieldOperator AST (FOAST) to Iterator IR (ITIR).
@@ -61,6 +65,8 @@ class FieldOperatorLowering(NodeTranslator):
     [Sym(id=SymbolName('inp'), kind='Iterator', dtype=('float64', False))]
     """
 
+    uid_generator: UIDGenerator = dataclasses.field(default_factory=UIDGenerator)
+
     @classmethod
     def apply(cls, node: foast.LocatedNode) -> itir.Expr:
         return cls().visit(node)
@@ -120,20 +126,78 @@ class FieldOperatorLowering(NodeTranslator):
         raise AssertionError("Statements must always be visited in the context of a function.")
 
     def visit_Return(
-        self, node: foast.Return, *, inner_expr: itir.Expr | None, **kwargs
+        self, node: foast.Return, *, inner_expr: Optional[itir.Expr], **kwargs
     ) -> itir.Expr:
         return self.visit(node.value, **kwargs)
 
     def visit_BlockStmt(
-        self, node: foast.BlockStmt, *, inner_expr: itir.Expr | None, **kwargs
+        self, node: foast.BlockStmt, *, inner_expr: Optional[itir.Expr], **kwargs
     ) -> itir.Expr:
         for stmt in reversed(node.stmts):
             inner_expr = self.visit(stmt, inner_expr=inner_expr, **kwargs)
         assert inner_expr
         return inner_expr
 
+    def visit_IfStmt(
+        self, node: foast.IfStmt, *, inner_expr: Optional[itir.Expr], **kwargs
+    ) -> itir.Expr:
+        # the lowered if call doesn't need to be lifted as the condition can only originate
+        #  from a scalar value (and not a field)
+        assert (
+            isinstance(node.condition.type, ts.ScalarType)
+            and node.condition.type.kind == ts.ScalarKind.BOOL
+        )
+
+        cond = self.visit(node.condition, **kwargs)
+
+        return_kind: StmtReturnKind = deduce_stmt_return_kind(node)
+
+        common_symbols: dict[str, foast.Symbol] = node.annex.propagated_symbols
+
+        if return_kind is StmtReturnKind.NO_RETURN:
+            # pack the common symbols into a tuple
+            common_symrefs = im.make_tuple(*(im.ref(sym) for sym in common_symbols.keys()))
+
+            # apply both branches and extract the common symbols through the prepared tuple
+            true_branch = self.visit(node.true_branch, inner_expr=common_symrefs, **kwargs)
+            false_branch = self.visit(node.false_branch, inner_expr=common_symrefs, **kwargs)
+
+            # unpack the common symbols' tuple for `inner_expr`
+            for i, sym in enumerate(common_symbols.keys()):
+                inner_expr = im.let(sym, im.tuple_get(i, im.ref("__if_stmt_result")))(inner_expr)
+
+            # here we assume neither branch returns
+            return im.let("__if_stmt_result", im.if_(im.deref(cond), true_branch, false_branch))(
+                inner_expr
+            )
+        elif return_kind is StmtReturnKind.CONDITIONAL_RETURN:
+            common_syms = tuple(im.sym(sym) for sym in common_symbols.keys())
+            common_symrefs = tuple(im.ref(sym) for sym in common_symbols.keys())
+
+            # wrap the inner expression in a lambda function. note that this increases the
+            # operation count if both branches are evaluated.
+            inner_expr_name = self.uid_generator.sequential_id(prefix="__inner_expr")
+            inner_expr_evaluator = im.lambda_(*common_syms)(inner_expr)
+            inner_expr = im.call(inner_expr_name)(*common_symrefs)
+
+            true_branch = self.visit(node.true_branch, inner_expr=inner_expr, **kwargs)
+            false_branch = self.visit(node.false_branch, inner_expr=inner_expr, **kwargs)
+
+            return im.let(inner_expr_name, inner_expr_evaluator)(
+                im.if_(im.deref(cond), true_branch, false_branch)
+            )
+
+        assert return_kind is StmtReturnKind.UNCONDITIONAL_RETURN
+
+        # note that we do not duplicate `inner_expr` here since if both branches
+        #  return, `inner_expr` is ignored.
+        true_branch = self.visit(node.true_branch, inner_expr=inner_expr, **kwargs)
+        false_branch = self.visit(node.false_branch, inner_expr=inner_expr, **kwargs)
+
+        return im.if_(im.deref(cond), true_branch, false_branch)
+
     def visit_Assign(
-        self, node: foast.Assign, *, inner_expr: itir.Expr | None, **kwargs
+        self, node: foast.Assign, *, inner_expr: Optional[itir.Expr], **kwargs
     ) -> itir.Expr:
         return im.let(self.visit(node.target, **kwargs), self.visit(node.value, **kwargs))(
             inner_expr
diff --git a/src/gt4py/next/ffront/func_to_foast.py b/src/gt4py/next/ffront/func_to_foast.py
index f63cfc3b..03f1187a 100644
--- a/src/gt4py/next/ffront/func_to_foast.py
+++ b/src/gt4py/next/ffront/func_to_foast.py
@@ -167,7 +167,7 @@ class FieldOperatorParser(DialectParser[foast.FunctionDefinition]):
 
         new_body = self._visit_stmts(node.body, self._make_loc(node), **kwargs)
 
-        if deduce_stmt_return_kind(new_body) != StmtReturnKind.UNCONDITIONAL_RETURN:
+        if deduce_stmt_return_kind(new_body) == StmtReturnKind.NO_RETURN:
             raise FieldOperatorSyntaxError.from_AST(
                 node, msg="Function must return a value, but no return statement was found."
             )
@@ -389,6 +389,15 @@ class FieldOperatorParser(DialectParser[foast.FunctionDefinition]):
             type=ts.DeferredType(constraint=ts.DataType),
         )
 
+    def visit_If(self, node: ast.If, **kwargs) -> foast.IfStmt:
+        loc = self._make_loc(node)
+        return foast.IfStmt(
+            condition=self.visit(node.test, **kwargs),
+            true_branch=self._visit_stmts(node.body, loc, **kwargs),
+            false_branch=self._visit_stmts(node.orelse, loc, **kwargs),
+            location=loc,
+        )
+
     def _visit_stmts(
         self, stmts: list[ast.stmt], location: eve.SourceLocation, **kwargs
     ) -> foast.BlockStmt:
diff --git a/src/gt4py/next/iterator/embedded.py b/src/gt4py/next/iterator/embedded.py
index c5befa58..fea6e2be 100644
--- a/src/gt4py/next/iterator/embedded.py
+++ b/src/gt4py/next/iterator/embedded.py
@@ -72,7 +72,9 @@ FieldAxis: TypeAlias = (
 )  # TODO Offset should be removed, is sometimes used for sparse dimensions
 TupleAxis: TypeAlias = type[None]
 Axis: TypeAlias = Union[FieldAxis, TupleAxis]
-Scalar: TypeAlias = SupportsInt | SupportsFloat | np.int32 | np.int64 | np.float32 | np.float64
+Scalar: TypeAlias = (
+    SupportsInt | SupportsFloat | np.int32 | np.int64 | np.float32 | np.float64 | np.bool_
+)
 
 
 class SparseTag(Tag):
@@ -198,7 +200,7 @@ class Column(np.lib.mixins.NDArrayOperatorsMixin):
 
     def __init__(self, kstart: int, data: np.ndarray | Scalar) -> None:
         self.kstart = kstart
-        assert isinstance(data, (np.ndarray, Scalar))  # type: ignore # mypy bug
+        assert isinstance(data, (np.ndarray, Scalar))  # type: ignore # mypy bug #11673
         column_range = column_range_cvar.get()
         self.data = data if isinstance(data, np.ndarray) else np.full(len(column_range), data)
 
@@ -470,7 +472,7 @@ def promote_scalars(val: CompositeOfScalarOrField):
     elif isinstance(val, LocatedField):
         return val
     val_type = infer_dtype_like_type(val)
-    if np.issubdtype(val_type, np.number):
+    if isinstance(val, Scalar):  # type: ignore # mypy bug
         return constant_field(val)
     else:
         raise ValueError(
diff --git a/src/gt4py/next/iterator/ir.py b/src/gt4py/next/iterator/ir.py
index 1a909927..54b06ca8 100644
--- a/src/gt4py/next/iterator/ir.py
+++ b/src/gt4py/next/iterator/ir.py
@@ -68,6 +68,11 @@ class Literal(Expr):
     value: str
     type: str  # noqa: A003
 
+    @datamodels.validator("type")
+    def _type_validator(self: datamodels.DataModelTP, attribute: datamodels.Attribute, value):
+        if value not in TYPEBUILTINS:
+            raise ValueError(f"{value} is not a valid builtin type.")
+
 
 class NoneLiteral(Expr):
     _none_literal: int = 0
diff --git a/src/gt4py/next/iterator/ir_makers.py b/src/gt4py/next/iterator/ir_makers.py
index c7dfe9d0..8fec9aaa 100644
--- a/src/gt4py/next/iterator/ir_makers.py
+++ b/src/gt4py/next/iterator/ir_makers.py
@@ -218,9 +218,9 @@ def make_tuple(*args):
     return call("make_tuple")(*args)
 
 
-def tuple_get(tuple_expr, index):
-    """Create a tuple_get FunCall, shorthand for ``call("tuple_get")(tuple_expr, index)``."""
-    return call("tuple_get")(tuple_expr, index)
+def tuple_get(index, tuple_expr):
+    """Create a tuple_get FunCall, shorthand for ``call("tuple_get")(index, tuple_expr)``."""
+    return call("tuple_get")(index, tuple_expr)
 
 
 def if_(cond, true_val, false_val):
diff --git a/src/gt4py/next/iterator/transforms/collapse_tuple.py b/src/gt4py/next/iterator/transforms/collapse_tuple.py
index df75ff11..efe0941a 100644
--- a/src/gt4py/next/iterator/transforms/collapse_tuple.py
+++ b/src/gt4py/next/iterator/transforms/collapse_tuple.py
@@ -29,25 +29,47 @@ def _get_tuple_size(node: ir.Node) -> int:
 
 @dataclass(frozen=True)
 class CollapseTuple(eve.NodeTranslator):
-    """Transform `make_tuple(tuple_get(0, t), tuple_get(1, t), ..., tuple_get(N-1,t))` -> t."""
+    """
+    Simplifies `make_tuple`, `tuple_get` calls.
+
+      - `make_tuple(tuple_get(0, t), tuple_get(1, t), ..., tuple_get(N-1,t))` -> `t`
+      - `tuple_get(i, make_tuple(e_0, e_1, ..., e_i, ..., e_N))` -> `e_i`
+    """
 
     ignore_tuple_size: bool
+    collapse_make_tuple_tuple_get: bool
+    collapse_tuple_get_make_tuple: bool
 
     @classmethod
-    def apply(cls, node: ir.Node, ignore_tuple_size: bool = False) -> ir.Node:
+    def apply(
+        cls,
+        node: ir.Node,
+        *,
+        ignore_tuple_size: bool = False,
+        # the following options are mostly for allowing separate testing of the modes
+        collapse_make_tuple_tuple_get: bool = True,
+        collapse_tuple_get_make_tuple: bool = True,
+    ) -> ir.Node:
         """
-        Transform `make_tuple(tuple_get(0, t), tuple_get(1, t), ..., tuple_get(N-1,t))` -> t.
+        Simplifies `make_tuple`, `tuple_get` calls.
 
         If `ignore_tuple_size`, apply the transformation even if length of the inner tuple
         is greater than the length of the outer tuple.
         """
-        return cls(ignore_tuple_size).visit(node)
+        return cls(
+            ignore_tuple_size, collapse_make_tuple_tuple_get, collapse_tuple_get_make_tuple
+        ).visit(node)
 
     def visit_FunCall(self, node: ir.FunCall, **kwargs) -> ir.Node:
-        if node.fun == ir.SymRef(id="make_tuple") and all(
-            isinstance(arg, ir.FunCall) and arg.fun == ir.SymRef(id="tuple_get")
-            for arg in node.args
+        if (
+            self.collapse_make_tuple_tuple_get
+            and node.fun == ir.SymRef(id="make_tuple")
+            and all(
+                isinstance(arg, ir.FunCall) and arg.fun == ir.SymRef(id="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)
             first_expr = node.args[0].args[1]
 
@@ -55,9 +77,24 @@ class CollapseTuple(eve.NodeTranslator):
                 assert isinstance(v, ir.FunCall)
                 assert isinstance(v.args[0], ir.Literal)
                 if not (int(v.args[0].value) == i and v.args[1] == first_expr):
+                    # tuple argument differs, just continue with the rest of the tree
                     return self.generic_visit(node)
 
             if self.ignore_tuple_size or _get_tuple_size(first_expr) == len(node.args):
                 return first_expr
-
+        if (
+            self.collapse_tuple_get_make_tuple
+            and node.fun == ir.SymRef(id="tuple_get")
+            and isinstance(node.args[1], ir.FunCall)
+            and node.args[1].fun == ir.SymRef(id="make_tuple")
+            and isinstance(node.args[0], ir.Literal)
+        ):
+            # `tuple_get(i, make_tuple(e_0, e_1, ..., e_i, ..., e_N))` -> `e_i`
+            assert node.args[0].type in ir.INTEGER_BUILTINS
+            make_tuple_call = node.args[1]
+            idx = int(node.args[0].value)
+            assert idx < len(
+                make_tuple_call.args
+            ), f"Index {idx} is out of bounds for tuple of size {len(make_tuple_call.args)}"
+            return node.args[1].args[idx]
         return self.generic_visit(node)
diff --git a/src/gt4py/next/iterator/type_inference.py b/src/gt4py/next/iterator/type_inference.py
index f39648d5..b609ea81 100644
--- a/src/gt4py/next/iterator/type_inference.py
+++ b/src/gt4py/next/iterator/type_inference.py
@@ -417,7 +417,12 @@ BUILTIN_TYPES: dict[str, Type] = {
         ),
         ret=Val_BOOL_T1,
     ),
-    "if_": FunctionType(args=Tuple.from_elems(Val_BOOL_T1, Val_T0_T1, Val_T0_T1), ret=Val_T0_T1),
+    "if_": FunctionType(
+        args=Tuple.from_elems(
+            Val_BOOL_T1, Val(kind=T2, dtype=T0, size=T1), Val(kind=T2, dtype=T0, size=T1)
+        ),
+        ret=Val(kind=T2, dtype=T0, size=T1),
+    ),
     "lift": FunctionType(
         args=Tuple.from_elems(
             FunctionType(


Test patch:
diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py
index 821953cd..d85da941 100644
--- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py
+++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py
@@ -460,7 +460,6 @@ def test_tuple_return_2(reduction_setup, fieldview_backend):
     assert np.allclose(ref, rs.out)
 
 
-@pytest.mark.xfail
 def test_tuple_with_local_field_in_reduction_shifted(reduction_setup, fieldview_backend):
     rs = reduction_setup
     V2EDim = rs.V2EDim
@@ -480,12 +479,7 @@ def test_tuple_with_local_field_in_reduction_shifted(reduction_setup, fieldview_
         edge_field: Field[[Edge], float64], vertex_field: Field[[Vertex], float64]
     ) -> Field[[Edge], float64]:
         tup = edge_field(V2E), vertex_field
-        # the shift inside the reduction fails as tup is a tuple of iterators
-        #  (as it contains a local field) which can not be shifted
         red = neighbor_sum(tup[0] + vertex_field, axis=V2EDim)
-        # even if the above is fixed we need to be careful with a subsequent
-        #  shift as the lifted lambda will contain tup as an argument which -
-        #  again - can not be shifted.
         return red(E2V[0])
 
     reduce_tuple_element(a, b, out=out, offset_provider=rs.offset_provider)
diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_scalar_if.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_scalar_if.py
new file mode 100644
index 00000000..08aceac2
--- /dev/null
+++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_scalar_if.py
@@ -0,0 +1,407 @@
+# -*- coding: utf-8 -*-
+# GT4Py - GridTools Framework
+#
+# Copyright (c) 2014-2023, ETH Zurich
+# All rights reserved.
+#
+# This file is part of the GT4Py project and the GridTools framework.
+# GT4Py is free software: you can redistribute it and/or modify it under
+# the terms of the GNU General Public License as published by the
+# Free Software Foundation, either version 3 of the License, or any later
+# version. See the LICENSE.txt file at the top-level directory of this
+# distribution for a copy of the license or check <https://www.gnu.org/licenses/>.
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+#
+from functools import reduce
+
+import numpy as np
+import pytest
+
+from gt4py.next.ffront.decorator import field_operator
+from gt4py.next.ffront.fbuiltins import Field, float64
+from gt4py.next.ffront.foast_passes.type_deduction import FieldOperatorTypeDeductionError
+from gt4py.next.iterator.embedded import index_field, np_as_located_field
+from gt4py.next.program_processors.runners import gtfn_cpu
+
+from next_tests.integration_tests.feature_tests import cases
+from next_tests.integration_tests.feature_tests.cases import (
+    E2V,
+    V2E,
+    E2VDim,
+    Edge,
+    IDim,
+    Ioff,
+    JDim,
+    Joff,
+    KDim,
+    Koff,
+    V2EDim,
+    Vertex,
+    cartesian_case,
+    no_default_backend,
+    unstructured_case,
+)
+from next_tests.integration_tests.feature_tests.ffront_tests.ffront_test_utils import (
+    Cell,
+    fieldview_backend,
+    reduction_setup,
+    size,
+)
+
+
+@pytest.mark.parametrize("condition", [True, False])
+def test_simple_if(condition, cartesian_case):
+    if cartesian_case.backend in [gtfn_cpu.run_gtfn, gtfn_cpu.run_gtfn_imperative]:
+        pytest.xfail("If-stmts are not supported yet.")
+
+    @field_operator
+    def simple_if(a: cases.IField, b: cases.IField, condition: bool) -> cases.IField:
+        if condition:
+            result = a
+        else:
+            result = b
+        return result
+
+    a = cases.allocate(cartesian_case, simple_if, "a")()
+    b = cases.allocate(cartesian_case, simple_if, "b")()
+    out = cases.allocate(cartesian_case, simple_if, cases.RETURN)()
+
+    cases.verify(cartesian_case, simple_if, a, b, condition, out=out, ref=a if condition else b)
+
+
+@pytest.mark.parametrize("condition1, condition2", [[True, False], [True, False]])
+def test_simple_if_conditional(condition1, condition2, cartesian_case):
+    if cartesian_case.backend in [gtfn_cpu.run_gtfn, gtfn_cpu.run_gtfn_imperative]:
+        pytest.xfail("If-stmts are not supported yet.")
+
+    @field_operator
+    def simple_if(
+        a: cases.IField,
+        b: cases.IField,
+        condition1: bool,
+        condition2: bool,
+    ) -> cases.IField:
+        if condition1:
+            result1 = a
+            result2 = a + 1
+        else:
+            result1 = b
+            result2 = b + 1
+        return result1 if condition2 else result2
+
+    a = cases.allocate(cartesian_case, simple_if, "a")()
+    b = cases.allocate(cartesian_case, simple_if, "b")()
+    out = cases.allocate(cartesian_case, simple_if, cases.RETURN)()
+
+    cases.verify(
+        cartesian_case,
+        simple_if,
+        a,
+        b,
+        condition1,
+        condition2,
+        out=out,
+        ref=(a if condition1 else b).array() + (0 if condition2 else 1),
+    )
+
+
+@pytest.mark.parametrize("condition", [True, False])
+def test_local_if(cartesian_case, condition):
+    if cartesian_case.backend in [gtfn_cpu.run_gtfn, gtfn_cpu.run_gtfn_imperative]:
+        pytest.xfail("If-stmts are not supported yet.")
+
+    @field_operator
+    def local_if(a: cases.IField, b: cases.IField, condition: bool) -> cases.IField:
+        if condition:
+            tmp = a
+            result = tmp
+        else:
+            result = b
+        return result
+
+    a = cases.allocate(cartesian_case, local_if, "a")()
+    b = cases.allocate(cartesian_case, local_if, "b")()
+    out = cases.allocate(cartesian_case, local_if, cases.RETURN)()
+
+    cases.verify(cartesian_case, local_if, a, b, condition, out=out, ref=(a if condition else b))
+
+
+@pytest.mark.parametrize("condition", [True, False])
+def test_temporary_if(cartesian_case, condition):
+    if cartesian_case.backend in [gtfn_cpu.run_gtfn, gtfn_cpu.run_gtfn_imperative]:
+        pytest.xfail("If-stmts are not supported yet.")
+
+    @field_operator
+    def temporary_if(a: cases.IField, b: cases.IField, condition: bool) -> cases.IField:
+        if condition:
+            tmp1 = a
+            result = tmp1
+        else:
+            tmp2 = b
+            result = tmp2
+        return result
+
+    a = cases.allocate(cartesian_case, temporary_if, "a")()
+    b = cases.allocate(cartesian_case, temporary_if, "b")()
+    out = cases.allocate(cartesian_case, temporary_if, cases.RETURN)()
+
+    cases.verify(
+        cartesian_case, temporary_if, a, b, condition, out=out, ref=(a if condition else b)
+    )
+
+
+@pytest.mark.parametrize("condition", [True, False])
+def test_if_return(cartesian_case, condition):
+    if cartesian_case.backend in [gtfn_cpu.run_gtfn, gtfn_cpu.run_gtfn_imperative]:
+        pytest.xfail("If-stmts are not supported yet.")
+
+    @field_operator
+    def temporary_if(a: cases.IField, b: cases.IField, condition: bool) -> cases.IField:
+        if condition:
+            tmp1 = a
+            return tmp1
+        else:
+            tmp2 = b
+            return tmp2
+        return a + b
+
+    a = cases.allocate(cartesian_case, temporary_if, "a")()
+    b = cases.allocate(cartesian_case, temporary_if, "b")()
+    out = cases.allocate(cartesian_case, temporary_if, cases.RETURN)()
+
+    cases.verify(
+        cartesian_case, temporary_if, a, b, condition, out=out, ref=(a if condition else b)
+    )
+
+
+@pytest.mark.parametrize("condition", [True, False])
+def test_if_stmt_if_branch_returns(cartesian_case, condition):
+    if cartesian_case.backend in [gtfn_cpu.run_gtfn, gtfn_cpu.run_gtfn_imperative]:
+        pytest.xfail("If-stmts are not supported yet.")
+
+    @field_operator
+    def if_branch_returns(a: cases.IField, b: cases.IField, condition: bool) -> cases.IField:
+        if condition:
+            tmp1 = a
+            return tmp1
+        return b
+
+    a = cases.allocate(cartesian_case, if_branch_returns, "a")()
+    b = cases.allocate(cartesian_case, if_branch_returns, "b")()
+    out = cases.allocate(cartesian_case, if_branch_returns, cases.RETURN)()
+
+    cases.verify(
+        cartesian_case, if_branch_returns, a, b, condition, out=out, ref=(a if condition else b)
+    )
+
+
+@pytest.mark.parametrize("condition", [True, False])
+def test_if_stmt_else_branch_returns(cartesian_case, condition):
+    if cartesian_case.backend in [gtfn_cpu.run_gtfn, gtfn_cpu.run_gtfn_imperative]:
+        pytest.xfail("If-stmts are not supported yet.")
+
+    @field_operator
+    def else_branch_returns(a: cases.IField, b: cases.IField, condition: bool) -> cases.IField:
+        if condition:
+            pass
+        else:
+            tmp1 = b
+            return tmp1
+        return a
+
+    a = cases.allocate(cartesian_case, else_branch_returns, "a")()
+    b = cases.allocate(cartesian_case, else_branch_returns, "b")()
+    out = cases.allocate(cartesian_case, else_branch_returns, cases.RETURN)()
+
+    cases.verify(
+        cartesian_case, else_branch_returns, a, b, condition, out=out, ref=(a if condition else b)
+    )
+
+
+@pytest.mark.parametrize("condition", [True, False])
+def test_if_stmt_both_branches_return(cartesian_case, condition):
+    if cartesian_case.backend in [gtfn_cpu.run_gtfn, gtfn_cpu.run_gtfn_imperative]:
+        pytest.xfail("If-stmts are not supported yet.")
+
+    @field_operator
+    def both_branches_return(a: cases.IField, b: cases.IField, condition: bool) -> cases.IField:
+        if condition:
+            tmp1 = a
+            return tmp1
+        else:
+            tmp2 = b
+            return tmp2
+
+    a = cases.allocate(cartesian_case, both_branches_return, "a")()
+    b = cases.allocate(cartesian_case, both_branches_return, "b")()
+    out = cases.allocate(cartesian_case, both_branches_return, cases.RETURN)()
+
+    cases.verify(
+        cartesian_case, both_branches_return, a, b, condition, out=out, ref=(a if condition else b)
+    )
+
+
+@pytest.mark.parametrize("condition1, condition2", [[True, False], [True, False]])
+def test_nested_if_stmt_conditinal(cartesian_case, condition1, condition2):
+    if cartesian_case.backend in [gtfn_cpu.run_gtfn, gtfn_cpu.run_gtfn_imperative]:
+        pytest.xfail("If-stmts are not supported yet.")
+
+    @field_operator
+    def nested_if_conditional_return(
+        inp: cases.IField, condition1: bool, condition2: bool
+    ) -> cases.IField:
+        if condition1:
+            tmp1 = inp
+            if condition2:
+                return tmp1 + 1
+            result = tmp1 + 2
+        else:
+            result = inp + 3
+        return result
+
+    inp = cases.allocate(cartesian_case, nested_if_conditional_return, "inp")()
+    out = cases.allocate(cartesian_case, nested_if_conditional_return, cases.RETURN)()
+
+    ref = {
+        (True, True): np.asarray(inp) + 1,
+        (True, False): np.asarray(inp) + 2,
+        (False, True): np.asarray(inp) + 3,
+        (False, False): np.asarray(inp) + 3,
+    }
+
+    cases.verify(
+        cartesian_case,
+        nested_if_conditional_return,
+        inp,
+        condition1,
+        condition2,
+        out=out,
+        ref=ref[(condition1, condition2)],
+    )
+
+
+@pytest.mark.parametrize("condition", [True, False])
+def test_nested_if(cartesian_case, condition):
+    if cartesian_case.backend in [gtfn_cpu.run_gtfn, gtfn_cpu.run_gtfn_imperative]:
+        pytest.xfail("If-stmts are not supported yet.")
+
+    @field_operator
+    def nested_if(a: cases.IField, b: cases.IField, condition: bool) -> cases.IField:
+        if condition:
+            if not condition:
+                inner = a
+            else:
+                inner = a + 1
+            result = inner
+        else:
+            result = b
+            if condition:
+                another_inner = 3
+            else:
+                another_inner = 5
+            result = result + another_inner
+        return result
+
+    a = cases.allocate(cartesian_case, nested_if, "a")()
+    b = cases.allocate(cartesian_case, nested_if, "b")()
+    out = cases.allocate(cartesian_case, nested_if, cases.RETURN)()
+
+    cases.verify(
+        cartesian_case,
+        nested_if,
+        a,
+        b,
+        condition,
+        out=out,
+        ref=np.asarray(a) + 1 if condition else np.asarray(b) + 5,
+    )
+
+
+@pytest.mark.parametrize("condition1, condition2", [[True, False], [True, False]])
+def test_if_without_else(cartesian_case, condition1, condition2):
+    if cartesian_case.backend in [gtfn_cpu.run_gtfn, gtfn_cpu.run_gtfn_imperative]:
+        pytest.xfail("If-stmts are not supported yet.")
+
+    @field_operator
+    def if_without_else(
+        a: cases.IField, b: cases.IField, condition1: bool, condition2: bool
+    ) -> cases.IField:
+        result = b + 1
+
+        if condition1:
+            if not condition2:
+                inner = a
+            else:
+                inner = a + 2
+            result = inner
+        return result
+
+    a = cases.allocate(cartesian_case, if_without_else, "a")()
+    b = cases.allocate(cartesian_case, if_without_else, "b")()
+    out = cases.allocate(cartesian_case, if_without_else, cases.RETURN)()
+
+    ref = {
+        (True, True): np.asarray(a) + 2,
+        (True, False): np.asarray(a),
+        (False, True): np.asarray(b) + 1,
+        (False, False): np.asarray(b) + 1,
+    }
+
+    cases.verify(
+        cartesian_case,
+        if_without_else,
+        a,
+        b,
+        condition1,
+        condition2,
+        out=out,
+        ref=ref[(condition1, condition2)],
+    )
+
+
+def test_if_non_scalar_condition():
+    with pytest.raises(FieldOperatorTypeDeductionError, match="Condition for `if` must be scalar."):
+
+        @field_operator
+        def if_non_scalar_condition(
+            a: Field[[IDim, JDim], float64], b: Field[[IDim, JDim], float64], condition: bool
+        ):
+            result = a
+            if a == b:
+                result = b
+            return result
+
+
+def test_if_non_boolean_condition():
+    with pytest.raises(
+        FieldOperatorTypeDeductionError, match="Condition for `if` must be of boolean type."
+    ):
+
+        @field_operator
+        def if_non_boolean_condition(
+            a: Field[[IDim, JDim], float64], b: Field[[IDim, JDim], float64], condition: float64
+        ):
+            if condition:
+                result = b
+            else:
+                result = a
+            return result
+
+
+def test_if_inconsistent_types():
+    with pytest.raises(
+        FieldOperatorTypeDeductionError,
+        match="Inconsistent types between two branches for variable",
+    ):
+
+        @field_operator
+        def if_inconsistent_types(
+            a: Field[[IDim, JDim], float64], b: Field[[IDim, JDim], float64], condition: bool
+        ):
+            if condition:
+                result = 1
+            else:
+                result = 2.0
+            return result
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 1dde3a68..f12e6b37 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
@@ -136,11 +136,11 @@ def test_eq():
 def test_if():
     testee = ir.SymRef(id="if_")
     c = ti.Val(kind=ti.Value(), dtype=ti.Primitive(name="bool"), size=ti.TypeVar(idx=0))
-    t = ti.Val(kind=ti.Value(), dtype=ti.TypeVar(idx=1), size=ti.TypeVar(idx=0))
+    t = ti.Val(kind=ti.TypeVar(idx=1), dtype=ti.TypeVar(idx=2), size=ti.TypeVar(idx=0))
     expected = ti.FunctionType(args=ti.Tuple.from_elems(c, t, t), ret=t)
     inferred = ti.infer(testee)
     assert inferred == expected
-    assert ti.pformat(inferred) == "(bool⁰, T₁⁰, T₁⁰) → T₁⁰"
+    assert ti.pformat(inferred) == "(bool⁰, ItOrVal₁[T₂⁰], ItOrVal₁[T₂⁰]) → ItOrVal₁[T₂⁰]"
 
 
 def test_not():
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 5860c0c6..736bf04d 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
@@ -14,47 +14,62 @@
 
 import pytest
 
-from gt4py.next.iterator import ir
+from gt4py.next.iterator import ir, ir_makers as im
 from gt4py.next.iterator.transforms.collapse_tuple import CollapseTuple
 
 
-def _tuple_get(i: int, t: ir.Expr):
-    return ir.FunCall(fun=ir.SymRef(id="tuple_get"), args=[ir.Literal(value=str(i), type="int"), t])
+def test_simple_make_tuple_tuple_get():
+    tuple_of_size_2 = im.make_tuple("first", "second")
+    testee = im.make_tuple(im.tuple_get(0, tuple_of_size_2), im.tuple_get(1, tuple_of_size_2))
 
+    actual = CollapseTuple.apply(testee, collapse_tuple_get_make_tuple=False)
 
-@pytest.fixture
-def tup_of_size_2():
-    return ir.FunCall(
-        fun=ir.SymRef(id="make_tuple"), args=[ir.SymRef(id="foo"), ir.SymRef(id="foo")]
-    )
+    expected = tuple_of_size_2
+    assert actual == expected
 
 
-def test_simple(tup_of_size_2):
-    testee = ir.FunCall(
-        fun=ir.SymRef(id="make_tuple"),
-        args=[_tuple_get(0, tup_of_size_2), _tuple_get(1, tup_of_size_2)],
+def test_nested_make_tuple_tuple_get():
+    tup_of_size2_from_lambda = im.call(im.lambda_()(im.make_tuple("first", "second")))()
+    testee = im.make_tuple(
+        im.tuple_get(0, tup_of_size2_from_lambda), im.tuple_get(1, tup_of_size2_from_lambda)
     )
-    expected = tup_of_size_2
-    actual = CollapseTuple.apply(testee)
-    assert actual == expected
 
+    actual = CollapseTuple.apply(testee, collapse_tuple_get_make_tuple=False)
 
-def test_incompatible_order(tup_of_size_2):
-    testee = ir.FunCall(
-        fun=ir.SymRef(id="make_tuple"),
-        args=[_tuple_get(1, tup_of_size_2), _tuple_get(0, tup_of_size_2)],
-    )
-    actual = CollapseTuple.apply(testee)
+    assert actual == tup_of_size2_from_lambda
+
+
+def test_different_tuples_make_tuple_tuple_get():
+    t0 = im.make_tuple("foo0", "bar0")
+    t1 = im.make_tuple("foo1", "bar1")
+    testee = im.make_tuple(im.tuple_get(0, t0), im.tuple_get(1, t1))
+
+    actual = CollapseTuple.apply(testee, collapse_tuple_get_make_tuple=False)
+
+    assert actual == testee  # did nothing
+
+
+def test_incompatible_order_make_tuple_tuple_get():
+    tuple_of_size_2 = im.make_tuple("first", "second")
+    testee = im.make_tuple(im.tuple_get(1, tuple_of_size_2), im.tuple_get(0, tuple_of_size_2))
+    actual = CollapseTuple.apply(testee, collapse_tuple_get_make_tuple=False)
     assert actual == testee  # did nothing
 
 
-def test_incompatible_size(tup_of_size_2):
-    testee = ir.FunCall(fun=ir.SymRef(id="make_tuple"), args=[_tuple_get(0, tup_of_size_2)])
-    actual = CollapseTuple.apply(testee)
+def test_incompatible_size_make_tuple_tuple_get():
+    testee = im.make_tuple(im.tuple_get(0, im.make_tuple("first", "second")))
+    actual = CollapseTuple.apply(testee, collapse_tuple_get_make_tuple=False)
     assert actual == testee  # did nothing
 
 
-def test_merged_with_smaller_outer_size(tup_of_size_2):
-    testee = ir.FunCall(fun=ir.SymRef(id="make_tuple"), args=[_tuple_get(0, tup_of_size_2)])
+def test_merged_with_smaller_outer_size_make_tuple_tuple_get():
+    testee = im.make_tuple(im.tuple_get(0, im.make_tuple("first", "second")))
     actual = CollapseTuple.apply(testee, ignore_tuple_size=True)
-    assert actual == tup_of_size_2
+    assert actual == im.make_tuple("first", "second")
+
+
+def test_simple_tuple_get_make_tuple():
+    expected = im.ref("bar")
+    testee = im.tuple_get(1, im.make_tuple("foo", expected))
+    actual = CollapseTuple.apply(testee, collapse_make_tuple_tuple_get=False)
+    assert expected == actual