Repository: GridTools/gt4py
Pull request #1058 (2022-11-22 21:07:16)

Fix can_deref in unroll_reduce
can_deref should be applied on a shifted iterator, not a sparse field.

TODO:
- [x] add test



Patch:
diff --git a/docs/functional/architecture/0013-Scalar_vs_0d_Fields.md b/docs/functional/architecture/0013-Scalar_vs_0d_Fields.md
new file mode 100644
index 00000000..56830a10
--- /dev/null
+++ b/docs/functional/architecture/0013-Scalar_vs_0d_Fields.md
@@ -0,0 +1,139 @@
+---
+tags: [frontend]
+---
+
+# Scalar vs 0d-fields
+
+- **Status**: valid
+- **Authors**: Hannes Vogt (@havogt)
+- **Created**: 2022-10-03
+- **Updated**: 2022-10-03
+
+Do we distinguish scalars (e.g. `float`) and fields with 0 dimensions (e.g. `Field[[], float]`) in the context of field view?
+
+## Context
+
+In the context of multi-dimensional fields, the question arises if we need scalars **and** 0-d fields, only one of them or even make them the same thing. Similar discussions happened within the Numpy community, see e.g. https://mail.python.org/pipermail/numpy-discussion/2006-February/006384.html.
+
+## Decision
+
+We decided to have both scalars and 0d-fields. They are not treated as alias, but have their distinct properties. Python typing semantics makes it difficult to have a nice solution and we want to comply with Python semantics for typing, especially for embedded execution.
+
+### Examples
+
+#### Wrong annotation for return type
+
+The following is forbidden because the return type annotation doesn't match the type of `1.0` (the correct annotation would be `float`).
+Note: the annotation (by Python design) is for the function that is being decorated, not for the function after the decorator is applied, see also this discussion https://github.com/python/typing/issues/412.
+
+```python=
+@field_operator
+def foo() -> Field[[], float]
+    return 1.0
+```
+
+An explicit broadcast would be required to make this correct.
+
+```python=
+@field_operator
+def foo() -> Field[[], float]
+    return broadcast(1.0, ())
+```
+
+#### Implicit broadcast to Field
+
+In the following call an implicit broadcast to Field is performed in the decorator (the annotation influences behavior).
+
+```python=
+@field_operator
+def foo(scalar: Field[[], float]) -> Field[[], float]:
+    return scalar
+
+foo(1.0)
+```
+
+Note that the following is also valid as no broadcasting to `Field` is applied.
+
+```python=
+@field_operator
+def foo(scalar: float) -> float:
+    return scalar
+
+foo(1.0)
+```
+
+This enables propagating scalars for cases where a scalars are required.
+
+#### 0-d fields cannot be used where a scalar is expected
+
+The following is forbidden because ternaries only work with scalar conditions.
+
+```python=
+@field_operator
+def foo(cond: Field[[], int]):
+    return some_field_op() if cond == 1 else some_other_field_op()
+```
+
+If needed, we might add functionality to extract the scalar value from a 0-d field.
+
+### Notes on lowering
+
+This decision doesn't change the current definition of iterator ir or the `fn` C++ backend:
+- a stencil (any lowered field_operator) will only accept iterators (i.e. `deref` is required);
+- `fn` will accept SIDs, scalars have to we wrapped in a helper, e.g. `global_parameter`.
+
+## Alternatives Considered
+
+### Scalar `T` and 0-d field `Field[[], T]` are aliases for each other
+
+Making scalars and 0-d field the same (aliases for each other) simplifies usage of the language in many places.
+
+Example: A reduction over a dimension of a field returns a field with one dimension less. Reducing over a 1-d field will therefore return a 0-d field, but that value might be needed in a place where a scalar is expected.
+
+However, we would need implement it with either possibly strange side-effects (making builtin types, e.g. `float` a subclass of Field) or create custom types for scalars (which is not desirable from a usage perspective, e.g. `gt4py.float`) to make this properly work with Python typing.
+
+It would allow to make some cases work that are forbidden by the now decided rules, e.g. 
+
+```python=
+@field_operator
+def foo() -> Field[[], float]
+    return 1.0
+```
+
+would be correct, e.g. if we implement `Field[[], float]` as a protocol that matches `1.0`.
+
+If we implement our own types, we have full freedom, to really make them aliases and get the following semantic work
+
+```python=
+scalar: gt4py.float(4.0)
+assert isinstance(scalar, gt4py.float)
+assert isinstance(scalar, Field[[], gt4py.float])
+assert scalar.dtype == gt4py.float
+assert scalar.dtype == Field[[], gt4py.float]
+assert scalar.dtype.dtype.dtype.dtype == gt4py.float
+assert scalar.dtype.dtype.dtype.dtype == Field[[], gt4py.float]
+```
+
+### 0-d fields are regular fields implementing scalar conversion operators
+
+In this proposal 0-d fields are regular fields which additionally implement the appropriate scalar conversion operator matching its datatype (`__int__` for integer types, `__float__` for floating point types, `__bool__` for booleans and `__complex__` for complex types). This would allow to use 0-d fields in places where Python require scalars (e.g. `if` conditions, sequence indexing).
+
+```python
+a_0d_field = field(2.3)
+
+assert isinstance(a_0d_field, Field)
+
+if a_0d_field > 0.5:
+    print("No problem!")
+```
+
+This is essentially the same approach used in the [Python array API standard](https://data-apis.org/array-api/latest/API_specification/array_object.html) where all methods dealing with arrays return arrays, and 0-d arrays support the conversion to scalars using standard _dunder_ methods. Additionally, this alternative is fully compatible with the implementation of our own custom scalar types, although it doesn't require it.
+
+In the long-term, this alternative seems to combine the best features from other options, but it requires a bit more discussion and a slightly larger implementation effort, so it has been postponed for now.
+
+## References <!-- optional -->
+
+- [\[Numpy-discussion\] A case for rank-0 arrays](https://mail.python.org/pipermail/numpy-discussion/2006-February/006384.html)
+- [Python - annotating the decorated function](https://github.com/python/typing/issues/412)
+- [Python array API standard - Array object](https://data-apis.org/array-api/latest/API_specification/array_object.html)
+- [GT4Py discussion and notes (HackMD)](https://hackmd.io/@gridtools/SyQ4vJ9Js)
diff --git a/docs/functional/architecture/Index.md b/docs/functional/architecture/Index.md
index 25fcc866..08dd5e2d 100644
--- a/docs/functional/architecture/Index.md
+++ b/docs/functional/architecture/Index.md
@@ -24,6 +24,7 @@ Writing a new ADR is simple:
 - [0001 - Field View Frontend Design](0001-Field_View_Frontend_Design.md)
 - [0002 - Field View Lowering](0002-Field_View_Lowering.md)
 - [0010 - Domain in Field View](0010-Domain_in_Field_View.md)
+- [0013 - Scalar vs 0d-Fields](0013-Scalar_vs_0d_Fields.md)
 
 ### Iterator IR #iterator
 - [0003 - Iterator View Tuple Support for Fields](0003-Iterator_View_Tuple_Support_for_Fields.md)
@@ -41,11 +42,15 @@ _None_
 - [0008 - Mapping Domain to Cpp Backend](0008-Mapping_Domain_to_Cpp-Backend.md)
 
 ### Python Integration
-- [0009 - Compiled Backend Integration](0009-Compiled-Backend-Integration.md)
+- [0011 - On The Fly Compilation](0011-On_The_Fly_Compilation.md)
+- [0012 - GridTools C++ OTF](0011-_GridTools_Cpp_OTF.md)
 
 ### Miscellanea
 _None_
 
+### Superseded
+- [0009 - Compiled Backend Integration](0009-Compiled-Backend-Integration.md)
+
 ## Other References
 
 - [GitHub Blog - Why Write ADRs. How architecture decision records can help your team](https://github.blog/2020-08-13-why-write-adrs/)
diff --git a/requirements-dev.txt b/requirements-dev.txt
index c338d719..a766227c 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -1,4 +1,5 @@
 check-manifest>=0.40
+clang-format>=15.0.6
 cmake>=3.18.1
 coverage>=5.0
 darglint>=1.6
diff --git a/setup.cfg b/setup.cfg
index 344d4402..a1f9a882 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -224,9 +224,6 @@ disallow_incomplete_defs = False
 [mypy-functional.program_processors.codegens/gtfn/gtfn_backend]
 ignore_errors = True
 
-[mypy-functional.program_processors.codegens/gtfn/itir_to_gtfn_ir]
-ignore_errors = True
-
 [mypy-functional.ffront.decorator]
 ignore_errors = True
 
@@ -242,7 +239,22 @@ ignore_errors = True
 [mypy-functional.iterator.runtime]
 ignore_errors = True
 
-[mypy-functional.iterator.transforms.*]
+[mypy-functional.iterator.transforms.global_tmps]
+ignore_errors = True
+
+[mypy-functional.iterator.transforms.remap_symbols]
+ignore_errors = True
+
+[mypy-functional.iterator.transforms.popup_tmps]
+ignore_errors = True
+
+[mypy-functional.iterator.transforms.merge_let]
+ignore_errors = True
+
+[mypy-functional.iterator.transforms.inline_lifts]
+ignore_errors = True
+
+[mypy-functional.iterator.transforms.inline_lambdas]
 ignore_errors = True
 
 #-- pytest --
diff --git a/src/functional/ffront/foast_passes/type_deduction.py b/src/functional/ffront/foast_passes/type_deduction.py
index 85f80007..d05268f5 100644
--- a/src/functional/ffront/foast_passes/type_deduction.py
+++ b/src/functional/ffront/foast_passes/type_deduction.py
@@ -400,7 +400,7 @@ class FieldOperatorTypeDeduction(traits.VisitorWithSymbolTableTrait, NodeTransla
         if condition.type != ts.ScalarType(kind=ts.ScalarKind.BOOL):
             raise FieldOperatorTypeDeductionError.from_foast_node(
                 condition,
-                msg=f"Condition is of type `{condition.type}` " f"but should be of type bool.",
+                msg=f"Condition is of type `{condition.type}` " f"but should be of type `bool`.",
             )
 
         if true_expr.type != false_expr.type:
diff --git a/src/functional/ffront/past_to_itir.py b/src/functional/ffront/past_to_itir.py
index e7aa9ee3..1fc47d6b 100644
--- a/src/functional/ffront/past_to_itir.py
+++ b/src/functional/ffront/past_to_itir.py
@@ -130,6 +130,7 @@ class ProgramLowering(traits.VisitorWithSymbolTableTrait, NodeTranslator):
         )
 
     def _visit_stencil_call(self, node: past.Call, **kwargs) -> itir.StencilClosure:
+        assert isinstance(node.kwargs["out"].type, ts.TypeSpec)
         assert type_info.is_field_type_or_tuple_of_field_type(node.kwargs["out"].type)
 
         domain = node.kwargs.get("domain", None)
@@ -193,6 +194,7 @@ class ProgramLowering(traits.VisitorWithSymbolTableTrait, NodeTranslator):
     ) -> itir.FunCall:
         domain_args = []
 
+        assert isinstance(out_field.type, ts.TypeSpec)
         out_field_types = type_info.primitive_constituents(out_field.type).to_list()
         out_field_types0_dims = cast(ts.FieldType, out_field_types[0]).dims
         if any(
diff --git a/src/functional/iterator/ir.py b/src/functional/iterator/ir.py
index 01b0b47a..9f6619c7 100644
--- a/src/functional/iterator/ir.py
+++ b/src/functional/iterator/ir.py
@@ -69,7 +69,7 @@ class FunctionDefinition(Node, SymbolTableTrait):
 
 
 class StencilClosure(Node):
-    domain: Expr
+    domain: FunCall
     stencil: Expr
     output: Union[SymRef, FunCall]
     inputs: List[SymRef]
diff --git a/src/functional/iterator/transforms/count_symbol_refs.py b/src/functional/iterator/transforms/count_symbol_refs.py
index a60476af..853ed0ad 100644
--- a/src/functional/iterator/transforms/count_symbol_refs.py
+++ b/src/functional/iterator/transforms/count_symbol_refs.py
@@ -9,7 +9,7 @@ class CountSymbolRefs(eve.NodeVisitor):
     ref_counts: dict[str, int]
 
     @classmethod
-    def apply(cls, node: itir.Node, symbol_names: list[str]) -> dict[str, int]:
+    def apply(cls, node: itir.Node | list[itir.Node], symbol_names: list[str]) -> dict[str, int]:
         ref_counts = {name: 0 for name in symbol_names}
         active_refs = set(symbol_names)
 
@@ -18,7 +18,7 @@ class CountSymbolRefs(eve.NodeVisitor):
 
         return obj.ref_counts
 
-    def visit_SymRef(self, node: itir.Node, *, active_refs: set[str]):
+    def visit_SymRef(self, node: itir.SymRef, *, active_refs: set[str]):
         if node.id in active_refs:
             self.ref_counts[node.id] += 1
 
diff --git a/src/functional/iterator/transforms/cse.py b/src/functional/iterator/transforms/cse.py
index 7b3ec9b5..1d9f8f3b 100644
--- a/src/functional/iterator/transforms/cse.py
+++ b/src/functional/iterator/transforms/cse.py
@@ -56,7 +56,7 @@ class CollectSubexpressions(NodeVisitor):
 
     def visit_FunCall(
         self,
-        node: ir.Lambda,
+        node: ir.FunCall,
         *,
         subexprs: dict[ir.Node, tuple[list[int], Optional[ir.Node]]],
         refs: ChainMap[str, bool],
diff --git a/src/functional/iterator/transforms/global_tmps.py b/src/functional/iterator/transforms/global_tmps.py
index 539785cb..3080ac20 100644
--- a/src/functional/iterator/transforms/global_tmps.py
+++ b/src/functional/iterator/transforms/global_tmps.py
@@ -40,7 +40,7 @@ Replaces lifted function calls by temporaries using the following steps:
 """
 
 
-AUTO_DOMAIN = ir.SymRef(id="_gtmp_auto_domain")
+AUTO_DOMAIN = ir.FunCall(fun=ir.SymRef(id="_gtmp_auto_domain"), args=[])
 
 
 # Iterator IR extension nodes
@@ -175,7 +175,9 @@ def split_closures(node: ir.FencilDefinition) -> FencilWithTemporaries:
         fencil=ir.FencilDefinition(
             id=node.id,
             function_definitions=node.function_definitions,
-            params=node.params + [ir.Sym(id=tmp.id) for tmp in tmps] + [ir.Sym(id=AUTO_DOMAIN.id)],
+            params=node.params
+            + [ir.Sym(id=tmp.id) for tmp in tmps]
+            + [ir.Sym(id=AUTO_DOMAIN.fun.id)],
             closures=list(reversed(closures)),
         ),
         params=node.params,
diff --git a/src/functional/iterator/transforms/inline_lambdas.py b/src/functional/iterator/transforms/inline_lambdas.py
index 30ded640..66668b15 100644
--- a/src/functional/iterator/transforms/inline_lambdas.py
+++ b/src/functional/iterator/transforms/inline_lambdas.py
@@ -13,6 +13,7 @@ def inline_lambda(
     force_inline_lift=False,
     eligible_params: Optional[list[bool]] = None,
 ):
+    assert isinstance(node.fun, ir.Lambda)
     eligible_params = eligible_params or [True] * len(node.fun.params)
 
     assert len(eligible_params) == len(node.fun.params) == len(node.args)
diff --git a/src/functional/iterator/transforms/simple_inline_heuristic.py b/src/functional/iterator/transforms/simple_inline_heuristic.py
index 05d105e2..ec0255fa 100644
--- a/src/functional/iterator/transforms/simple_inline_heuristic.py
+++ b/src/functional/iterator/transforms/simple_inline_heuristic.py
@@ -22,6 +22,7 @@ def is_eligible_for_inlining(node: ir.FunCall, is_scan_pass_context: bool) -> bo
     """
     assert _is_lift(node)
 
+    assert isinstance(node.fun, ir.FunCall)  # for mypy
     (stencil,) = node.fun.args
     # Don’t inline scans, i.e. exclude `↑(scan(...))(...)`
     if isinstance(stencil, ir.FunCall) and stencil.fun == ir.SymRef(id="scan"):
diff --git a/src/functional/iterator/transforms/unroll_reduce.py b/src/functional/iterator/transforms/unroll_reduce.py
index 877fc1f3..c2b6e142 100644
--- a/src/functional/iterator/transforms/unroll_reduce.py
+++ b/src/functional/iterator/transforms/unroll_reduce.py
@@ -1,11 +1,98 @@
 import dataclasses
-from collections.abc import Iterable
+from collections.abc import Iterable, Iterator
+from typing import TypeGuard
 
 from eve import NodeTranslator
 from eve.utils import UIDGenerator
+from functional import common
 from functional.iterator import ir
 
 
+def _is_shifted(arg: ir.Expr) -> TypeGuard[ir.FunCall]:
+    return (
+        isinstance(arg, ir.FunCall)
+        and isinstance(arg.fun, ir.FunCall)
+        and arg.fun.fun == ir.SymRef(id="shift")
+    )
+
+
+def _is_applied_lift(arg: ir.Expr) -> TypeGuard[ir.FunCall]:
+    return (
+        isinstance(arg, ir.FunCall)
+        and isinstance(arg.fun, ir.FunCall)
+        and arg.fun.fun == ir.SymRef(id="lift")
+    )
+
+
+def _is_shifted_or_lifted_and_shifted(arg: ir.Expr) -> TypeGuard[ir.FunCall]:
+    return _is_shifted(arg) or (
+        _is_applied_lift(arg)
+        and any(_is_shifted_or_lifted_and_shifted(nested_arg) for nested_arg in arg.args)
+    )
+
+
+def _get_shifted_args(reduce_args: Iterable[ir.Expr]) -> Iterator[ir.FunCall]:
+    return filter(
+        _is_shifted_or_lifted_and_shifted,
+        reduce_args,
+    )
+
+
+def _is_list_of_funcalls(lst: list) -> TypeGuard[list[ir.FunCall]]:
+    return all(isinstance(f, ir.FunCall) for f in lst)
+
+
+def _get_partial_offset(arg: ir.FunCall) -> ir.OffsetLiteral:
+    if _is_shifted(arg):
+        assert isinstance(arg.fun, ir.FunCall)
+        offset = arg.fun.args[-1]
+        assert isinstance(offset, ir.OffsetLiteral)
+        return offset
+    else:
+        assert _is_applied_lift(arg)
+        assert _is_list_of_funcalls(arg.args)
+        partial_offsets = [_get_partial_offset(arg) for arg in arg.args]
+        assert all(o == partial_offsets[0] for o in partial_offsets)
+        return partial_offsets[0]
+
+
+def _get_connectivity(reduce_args: Iterable[ir.Expr], offset_provider) -> common.Connectivity:
+    connectivities = []
+    for arg in _get_shifted_args(reduce_args):
+        connectivities.append(offset_provider[_get_partial_offset(arg).value])
+
+    if not connectivities:
+        raise RuntimeError("Couldn't detect partial shift in any arguments of reduce.")
+
+    if len({(c.max_neighbors, c.has_skip_values) for c in connectivities}) != 1:
+        # The condition for this check is required but not sufficient: the actual neighbor tables could still be incompatible.
+        raise RuntimeError("Arguments to reduce have incompatible partial shifts.")
+    return connectivities[0]
+
+
+def _is_reduce(node: ir.FunCall):
+    return isinstance(node.fun, ir.FunCall) and node.fun.fun == ir.SymRef(id="reduce")
+
+
+def _make_shift(offsets: list[ir.Expr], iterator: ir.Expr):
+    return ir.FunCall(fun=ir.FunCall(fun=ir.SymRef(id="shift"), args=offsets), args=[iterator])
+
+
+def _make_deref(iterator: ir.Expr):
+    return ir.FunCall(fun=ir.SymRef(id="deref"), args=[iterator])
+
+
+def _make_can_deref(iterator: ir.Expr):
+    return ir.FunCall(fun=ir.SymRef(id="can_deref"), args=[iterator])
+
+
+def _make_if(cond: ir.Expr, true_expr: ir.Expr, false_expr: ir.Expr):
+    return ir.FunCall(
+        fun=ir.SymRef(id="if_"),
+        args=[cond, true_expr, false_expr],
+    )
+
+
 @dataclasses.dataclass(frozen=True)
 class UnrollReduce(NodeTranslator):
     # we use one UID generator per instance such that the generated ids are
@@ -16,70 +103,14 @@ class UnrollReduce(NodeTranslator):
     def apply(cls, node: ir.Node, **kwargs):
         return cls().visit(node, **kwargs)
 
-    @staticmethod
-    def _find_connectivities(reduce_args: Iterable[ir.Expr], offset_provider):
-        connectivities = []
-        for arg in reduce_args:
-            if (
-                isinstance(arg, ir.FunCall)
-                and isinstance(arg.fun, ir.FunCall)
-                and arg.fun.fun == ir.SymRef(id="shift")
-            ):
-                assert isinstance(arg.fun.args[-1], ir.OffsetLiteral), f"{arg.fun.args}"
-                connectivities.append(offset_provider[arg.fun.args[-1].value])
-            # if the argument is a lift call by itself look at its arguments recursively
-            elif (
-                isinstance(arg, ir.FunCall)
-                and isinstance(arg.fun, ir.FunCall)
-                and arg.fun.fun == ir.SymRef(id="lift")
-            ):
-                connectivities.extend(UnrollReduce._find_connectivities(arg.args, offset_provider))
-
-        return connectivities
-
-    @staticmethod
-    def _find_connectivity(reduce_args: Iterable[ir.Expr], offset_provider):
-        connectivities = UnrollReduce._find_connectivities(reduce_args, offset_provider)
-
-        if not connectivities:
-            raise RuntimeError("Couldn't detect partial shift in any arguments of reduce.")
-
-        if len({(c.max_neighbors, c.has_skip_values) for c in connectivities}) != 1:
-            # The condition for this check is required but not sufficient: the actual neighbor tables could still be incompatible.
-            raise RuntimeError("Arguments to reduce have incompatible partial shifts.")
-        return connectivities[0]
-
-    @staticmethod
-    def _is_reduce(node: ir.FunCall):
-        return isinstance(node.fun, ir.FunCall) and node.fun.fun == ir.SymRef(id="reduce")
-
-    @staticmethod
-    def _make_shift(offsets: list[ir.Expr], iterator: ir.Expr):
-        return ir.FunCall(fun=ir.FunCall(fun=ir.SymRef(id="shift"), args=offsets), args=[iterator])
-
-    @staticmethod
-    def _make_deref(iterator: ir.Expr):
-        return ir.FunCall(fun=ir.SymRef(id="deref"), args=[iterator])
-
-    @staticmethod
-    def _make_can_deref(iterator: ir.Expr):
-        return ir.FunCall(fun=ir.SymRef(id="can_deref"), args=[iterator])
-
-    @staticmethod
-    def _make_if(cond: ir.Expr, true_expr: ir.Expr, false_expr: ir.Expr):
-        return ir.FunCall(
-            fun=ir.SymRef(id="if_"),
-            args=[cond, true_expr, false_expr],
-        )
-
     def visit_FunCall(self, node: ir.FunCall, **kwargs):
         node = self.generic_visit(node, **kwargs)
-        if not self._is_reduce(node):
+        if not _is_reduce(node):
             return node
 
         offset_provider = kwargs["offset_provider"]
         assert offset_provider is not None
-        connectivity = self._find_connectivity(node.args, offset_provider)
+        connectivity = _get_connectivity(node.args, offset_provider)
         max_neighbors = connectivity.max_neighbors
         has_skip_values = connectivity.has_skip_values
 
@@ -90,13 +121,12 @@ class UnrollReduce(NodeTranslator):
         assert isinstance(node.fun, ir.FunCall)
         fun, init = node.fun.args
 
-        derefed_shifted_args = [
-            self._make_deref(self._make_shift([offset], arg)) for arg in node.args
-        ]
+        derefed_shifted_args = [_make_deref(_make_shift([offset], arg)) for arg in node.args]
         step_fun: ir.Expr = ir.FunCall(fun=fun, args=[acc] + derefed_shifted_args)
         if has_skip_values:
-            can_deref = self._make_can_deref(self._make_shift([offset], node.args[0]))
-            step_fun = self._make_if(can_deref, step_fun, acc)
+            check_arg = next(_get_shifted_args(node.args))
+            can_deref = _make_can_deref(_make_shift([offset], check_arg))
+            step_fun = _make_if(can_deref, step_fun, acc)
         step_fun = ir.Lambda(params=[ir.Sym(id=acc.id), ir.Sym(id=offset.id)], expr=step_fun)
         expr = init
         for i in range(max_neighbors):
diff --git a/src/functional/program_processors/codegens/gtfn/itir_to_gtfn_ir.py b/src/functional/program_processors/codegens/gtfn/itir_to_gtfn_ir.py
index 143d329d..1fab4e47 100644
--- a/src/functional/program_processors/codegens/gtfn/itir_to_gtfn_ir.py
+++ b/src/functional/program_processors/codegens/gtfn/itir_to_gtfn_ir.py
@@ -70,9 +70,7 @@ _horizontal_dimension = "gtfn::unstructured::dim::horizontal"
 
 
 def _get_domains(closures: Iterable[itir.StencilClosure]) -> Iterable[itir.FunCall]:
-    domains = [c.domain for c in closures]
-    assert all(isinstance(d, itir.FunCall) for d in domains)
-    return domains
+    return (c.domain for c in closures)
 
 
 def _extract_grid_type(domain: itir.FunCall) -> common.GridType:
@@ -97,6 +95,7 @@ def _name_from_named_range(named_range_call: itir.FunCall) -> str:
     assert isinstance(named_range_call, itir.FunCall) and named_range_call.fun == itir.SymRef(
         id="named_range"
     )
+    assert isinstance(named_range_call.args[0], itir.AxisLiteral)
     return named_range_call.args[0].value
 
 
@@ -108,6 +107,7 @@ def _collect_dimensions_from_domain(
     for domain in domains:
         if domain.fun == itir.SymRef(id="cartesian_domain"):
             for nr in domain.args:
+                assert isinstance(nr, itir.FunCall)
                 dim_name = _name_from_named_range(nr)
                 offset_definitions[dim_name] = TagDefinition(name=Sym(id=dim_name))
         elif domain.fun == itir.SymRef(id="unstructured_domain"):
@@ -115,12 +115,14 @@ def _collect_dimensions_from_domain(
                 raise ValueError("unstructured_domain must not have more than 2 arguments.")
             if len(domain.args) > 0:
                 horizontal_range = domain.args[0]
+                assert isinstance(horizontal_range, itir.FunCall)
                 horizontal_name = _name_from_named_range(horizontal_range)
                 offset_definitions[horizontal_name] = TagDefinition(
                     name=Sym(id=horizontal_name), alias=_horizontal_dimension
                 )
             if len(domain.args) > 1:
                 vertical_range = domain.args[1]
+                assert isinstance(vertical_range, itir.FunCall)
                 vertical_name = _name_from_named_range(vertical_range)
                 offset_definitions[vertical_name] = TagDefinition(
                     name=Sym(id=vertical_name), alias=_vertical_dimension
@@ -149,17 +151,20 @@ def _collect_offset_definitions(
             raise ValueError(f"Missing offset_provider entry for {o.value}")
 
         offset_name = o.value
-        if isinstance(offset_provider[o.value], common.Dimension):
-            dim = offset_provider[o.value]
+        assert isinstance(offset_name, str)
+        dim_or_conn = offset_provider[offset_name]
+        if isinstance(dim_or_conn, common.Dimension):
             if grid_type == common.GridType.CARTESIAN:
                 # create alias from offset to dimension
-                offset_definitions[dim.value] = TagDefinition(name=Sym(id=dim.value))
+                offset_definitions[dim_or_conn.value] = TagDefinition(
+                    name=Sym(id=dim_or_conn.value)
+                )
                 offset_definitions[offset_name] = TagDefinition(
-                    name=Sym(id=offset_name), alias=SymRef(id=dim.value)
+                    name=Sym(id=offset_name), alias=SymRef(id=dim_or_conn.value)
                 )
             else:
                 assert grid_type == common.GridType.UNSTRUCTURED
-                if not dim.kind == common.DimensionKind.VERTICAL:
+                if not dim_or_conn.kind == common.DimensionKind.VERTICAL:
                     raise ValueError(
                         "Mapping an offset to a horizontal dimension in unstructured is not allowed."
                     )
@@ -168,7 +173,7 @@ def _collect_offset_definitions(
                     name=Sym(id=offset_name), alias=_vertical_dimension
                 )
         else:
-            assert isinstance(offset_provider[o.value], common.Connectivity)
+            assert isinstance(dim_or_conn, common.Connectivity)
             offset_definitions[offset_name] = TagDefinition(name=Sym(id=offset_name))
     return offset_definitions
 
@@ -340,6 +345,7 @@ class GTFN_lowering(eve.NodeTranslator, eve.VisitorWithSymbolTableTrait):
         return result
 
     def _error_on_illegal_function_calls(self, node: itir.FunCall) -> None:
+        assert isinstance(node.fun, itir.SymRef)
         if node.fun.id == "shift":
             raise ValueError("unapplied shift call not supported: {node}")
         elif node.fun.id == "scan":
@@ -419,7 +425,7 @@ class GTFN_lowering(eve.NodeTranslator, eve.VisitorWithSymbolTableTrait):
     def _is_scan(node: itir.Node):
         return isinstance(node, itir.FunCall) and node.fun == itir.SymRef(id="scan")
 
-    def _visit_output_argument(self, node: itir.SymRef | itir.FunCall):
+    def _visit_output_argument(self, node: itir.Expr):
         if isinstance(node, itir.SymRef):
             return self.visit(node)
         elif isinstance(node, itir.FunCall) and node.fun == itir.SymRef(id="make_tuple"):
diff --git a/src/functional/type_system/type_info.py b/src/functional/type_system/type_info.py
index 514b02c7..4befca47 100644
--- a/src/functional/type_system/type_info.py
+++ b/src/functional/type_system/type_info.py
@@ -14,7 +14,7 @@
 
 import functools
 from types import EllipsisType
-from typing import Callable, Iterator, Optional, Type, TypeGuard, cast
+from typing import Callable, Iterator, Type, TypeGuard, cast
 
 from eve.utils import XIterable, xiter
 from functional.common import Dimension, GTTypeError
@@ -60,7 +60,7 @@ def type_class(symbol_type: ts.TypeSpec) -> Type[ts.TypeSpec]:
 
 
 def primitive_constituents(
-    symbol_type: Optional[ts.TypeSpec],
+    symbol_type: ts.TypeSpec,
 ) -> XIterable[ts.TypeSpec]:
     """
     Return the primitive types contained in a composite type.
@@ -79,7 +79,7 @@ def primitive_constituents(
     [FieldType(...), ScalarType(...), FieldType(...)]
     """
 
-    def constituents_yielder(symbol_type: Optional[ts.TypeSpec]):
+    def constituents_yielder(symbol_type: ts.TypeSpec):
         if isinstance(symbol_type, ts.TupleType):
             for el_type in symbol_type.types:
                 yield from constituents_yielder(el_type)
@@ -228,7 +228,7 @@ def is_arithmetic(symbol_type: ts.TypeSpec) -> bool:
     return is_floating_point(symbol_type) or is_integral(symbol_type)
 
 
-def is_field_type_or_tuple_of_field_type(type_: Optional[ts.TypeSpec]) -> bool:
+def is_field_type_or_tuple_of_field_type(type_: ts.TypeSpec) -> bool:
     """
      Return True if ``type_`` is FieldType or FieldType nested in TupleType.
 


Test patch:
diff --git a/tests/functional_tests/ffront_tests/test_foast_parsing.py b/tests/functional_tests/ffront_tests/test_foast_parsing.py
index 48482aa2..39e97c66 100644
--- a/tests/functional_tests/ffront_tests/test_foast_parsing.py
+++ b/tests/functional_tests/ffront_tests/test_foast_parsing.py
@@ -28,7 +28,16 @@ import pytest
 from eve.pattern_matching import ObjectPattern as P
 from functional.common import Field, GTTypeError
 from functional.ffront import field_operator_ast as foast, type_specifications as ts
-from functional.ffront.fbuiltins import Dimension, astype, float32, float64, int32, int64, where
+from functional.ffront.fbuiltins import (
+    Dimension,
+    astype,
+    broadcast,
+    float32,
+    float64,
+    int32,
+    int64,
+    where,
+)
 from functional.ffront.foast_passes.type_deduction import FieldOperatorTypeDeductionError
 from functional.ffront.func_to_foast import FieldOperatorParser, FieldOperatorSyntaxError
 from functional.iterator import ir as itir
@@ -266,6 +275,54 @@ def test_conditional_wrong_arg_type():
     assert re.search(msg, exc_info.value.__cause__.args[0]) is not None
 
 
+def test_ternary_with_field_condition():
+    def ternary_with_field_condition(cond: Field[[], bool]):
+        return 1 if cond else 2
+
+    with pytest.raises(FieldOperatorTypeDeductionError, match=r"should be .* `bool`"):
+        _ = FieldOperatorParser.apply_to_function(ternary_with_field_condition)
+
+
+def test_correct_return_type_annotation():
+    """See ADR 13."""
+
+    def correct_return_type_annotation() -> float:
+        return 1.0
+
+    FieldOperatorParser.apply_to_function(correct_return_type_annotation)
+
+
+def test_adr13_wrong_return_type_annotation():
+    """See ADR 13."""
+
+    def wrong_return_type_annotation() -> Field[[], float]:
+        return 1.0
+
+    with pytest.raises(GTTypeError, match=r"Expected `float.*`"):
+        _ = FieldOperatorParser.apply_to_function(wrong_return_type_annotation)
+
+
+def test_adr13_fixed_return_type_annotation():
+    """See ADR 13."""
+
+    def fixed_return_type_annotation() -> Field[[], float]:
+        return broadcast(1.0, ())
+
+    FieldOperatorParser.apply_to_function(fixed_return_type_annotation)
+
+
+def test_no_implicit_broadcast_in_field_op_call():
+    """See ADR 13."""
+
+    def no_implicit_broadcast_in_field_op_call(scalar: float) -> float:
+        return scalar
+
+    def no_implicit_broadcast_in_field_op_call_caller() -> float:
+        return no_implicit_broadcast_in_field_op_call(1.0)
+
+    FieldOperatorParser.apply_to_function(no_implicit_broadcast_in_field_op_call_caller)
+
+
 def test_astype():
     def astype_fieldop(a: Field[..., "int64"]) -> Field[..., float64]:
         return astype(a, float64)
diff --git a/tests/functional_tests/iterator_tests/test_global_tmps.py b/tests/functional_tests/iterator_tests/test_global_tmps.py
index 07e2654f..0e868304 100644
--- a/tests/functional_tests/iterator_tests/test_global_tmps.py
+++ b/tests/functional_tests/iterator_tests/test_global_tmps.py
@@ -2,6 +2,7 @@ from eve.utils import UIDs
 from functional.iterator import ir
 from functional.iterator.runtime import CartesianAxis
 from functional.iterator.transforms.global_tmps import (
+    AUTO_DOMAIN,
     FencilWithTemporaries,
     Temporary,
     collect_tmps_info,
@@ -18,7 +19,7 @@ def test_split_closures():
         params=[ir.Sym(id="d"), ir.Sym(id="inp"), ir.Sym(id="out")],
         closures=[
             ir.StencilClosure(
-                domain=ir.SymRef(id="d"),
+                domain=ir.FunCall(fun=ir.SymRef(id="cartesian_domain"), args=[]),
                 stencil=ir.Lambda(
                     params=[ir.Sym(id="baz_inp")],
                     expr=ir.FunCall(
@@ -79,7 +80,7 @@ def test_split_closures():
         ],
         closures=[
             ir.StencilClosure(
-                domain=ir.SymRef(id="_gtmp_auto_domain"),
+                domain=AUTO_DOMAIN,
                 stencil=ir.Lambda(
                     params=[ir.Sym(id="foo_inp")],
                     expr=ir.FunCall(
@@ -91,7 +92,7 @@ def test_split_closures():
                 inputs=[ir.SymRef(id="inp")],
             ),
             ir.StencilClosure(
-                domain=ir.SymRef(id="_gtmp_auto_domain"),
+                domain=AUTO_DOMAIN,
                 stencil=ir.Lambda(
                     params=[
                         ir.Sym(id="bar_inp"),
@@ -108,7 +109,7 @@ def test_split_closures():
                 inputs=[ir.SymRef(id="inp"), ir.SymRef(id="_gtmp_1")],
             ),
             ir.StencilClosure(
-                domain=ir.SymRef(id="d"),
+                domain=ir.FunCall(fun=ir.SymRef(id="cartesian_domain"), args=[]),
                 stencil=ir.Lambda(
                     params=[ir.Sym(id="baz_inp"), ir.Sym(id="_lift_2")],
                     expr=ir.FunCall(
@@ -143,7 +144,7 @@ def test_update_cartesian_domains():
             ],
             closures=[
                 ir.StencilClosure(
-                    domain=ir.SymRef(id="_gtmp_auto_domain"),
+                    domain=AUTO_DOMAIN,
                     stencil=ir.Lambda(
                         params=[ir.Sym(id="foo_inp")],
                         expr=ir.FunCall(
@@ -155,7 +156,7 @@ def test_update_cartesian_domains():
                     inputs=[ir.SymRef(id="inp")],
                 ),
                 ir.StencilClosure(
-                    domain=ir.SymRef(id="_gtmp_auto_domain"),
+                    domain=AUTO_DOMAIN,
                     stencil=ir.SymRef(id="deref"),
                     output=ir.SymRef(id="_gtmp_0"),
                     inputs=[ir.SymRef(id="_gtmp_1")],
diff --git a/tests/functional_tests/iterator_tests/test_pretty_parser.py b/tests/functional_tests/iterator_tests/test_pretty_parser.py
index 5e4c60b9..cc5ea6e7 100644
--- a/tests/functional_tests/iterator_tests/test_pretty_parser.py
+++ b/tests/functional_tests/iterator_tests/test_pretty_parser.py
@@ -174,9 +174,9 @@ def test_function_definition():
 
 
 def test_stencil_closure():
-    testee = "y ← (deref)(x) @ d;"
+    testee = "y ← (deref)(x) @ cartesian_domain();"
     expected = ir.StencilClosure(
-        domain=ir.SymRef(id="d"),
+        domain=ir.FunCall(fun=ir.SymRef(id="cartesian_domain"), args=[]),
         stencil=ir.SymRef(id="deref"),
         output=ir.SymRef(id="y"),
         inputs=[ir.SymRef(id="x")],
@@ -186,7 +186,7 @@ def test_stencil_closure():
 
 
 def test_fencil_definition():
-    testee = "f(d, x, y) {\n  g = λ(x) → x;\n  y ← (deref)(x) @ d;\n}"
+    testee = "f(d, x, y) {\n  g = λ(x) → x;\n  y ← (deref)(x) @ cartesian_domain();\n}"
     expected = ir.FencilDefinition(
         id="f",
         function_definitions=[
@@ -195,7 +195,7 @@ def test_fencil_definition():
         params=[ir.Sym(id="d"), ir.Sym(id="x"), ir.Sym(id="y")],
         closures=[
             ir.StencilClosure(
-                domain=ir.SymRef(id="d"),
+                domain=ir.FunCall(fun=ir.SymRef(id="cartesian_domain"), args=[]),
                 stencil=ir.SymRef(id="deref"),
                 output=ir.SymRef(id="y"),
                 inputs=[ir.SymRef(id="x")],
diff --git a/tests/functional_tests/iterator_tests/test_pretty_printer.py b/tests/functional_tests/iterator_tests/test_pretty_printer.py
index fcbdb55f..4f8c8a92 100644
--- a/tests/functional_tests/iterator_tests/test_pretty_printer.py
+++ b/tests/functional_tests/iterator_tests/test_pretty_printer.py
@@ -295,12 +295,12 @@ def test_function_definition():
 
 def test_stencil_closure():
     testee = ir.StencilClosure(
-        domain=ir.SymRef(id="d"),
+        domain=ir.FunCall(fun=ir.SymRef(id="cartesian_domain"), args=[]),
         stencil=ir.SymRef(id="deref"),
         output=ir.SymRef(id="y"),
         inputs=[ir.SymRef(id="x")],
     )
-    expected = "y ← (deref)(x) @ d;"
+    expected = "y ← (deref)(x) @ cartesian_domain();"
     actual = pformat(testee)
     assert actual == expected
 
@@ -314,7 +314,7 @@ def test_fencil_definition():
         params=[ir.Sym(id="d"), ir.Sym(id="x"), ir.Sym(id="y")],
         closures=[
             ir.StencilClosure(
-                domain=ir.SymRef(id="d"),
+                domain=ir.FunCall(fun=ir.SymRef(id="cartesian_domain"), args=[]),
                 stencil=ir.SymRef(id="deref"),
                 output=ir.SymRef(id="y"),
                 inputs=[ir.SymRef(id="x")],
@@ -322,5 +322,5 @@ def test_fencil_definition():
         ],
     )
     actual = pformat(testee)
-    expected = "f(d, x, y) {\n  g = λ(x) → x;\n  y ← (deref)(x) @ d;\n}"
+    expected = "f(d, x, y) {\n  g = λ(x) → x;\n  y ← (deref)(x) @ cartesian_domain();\n}"
     assert actual == expected
diff --git a/tests/functional_tests/iterator_tests/test_prune_closure_inputs.py b/tests/functional_tests/iterator_tests/test_prune_closure_inputs.py
index 11b64d27..066eb6b9 100644
--- a/tests/functional_tests/iterator_tests/test_prune_closure_inputs.py
+++ b/tests/functional_tests/iterator_tests/test_prune_closure_inputs.py
@@ -4,7 +4,7 @@ from functional.iterator.transforms.prune_closure_inputs import PruneClosureInpu
 
 def test_simple():
     testee = ir.StencilClosure(
-        domain=ir.SymRef(id="d"),
+        domain=ir.FunCall(fun=ir.SymRef(id="cartesian_domain"), args=[]),
         stencil=ir.Lambda(
             params=[ir.Sym(id="x"), ir.Sym(id="y"), ir.Sym(id="z")],
             expr=ir.FunCall(fun=ir.SymRef(id="deref"), args=[ir.SymRef(id="y")]),
@@ -13,7 +13,7 @@ def test_simple():
         inputs=[ir.SymRef(id="foo"), ir.SymRef(id="bar"), ir.SymRef(id="baz")],
     )
     expected = ir.StencilClosure(
-        domain=ir.SymRef(id="d"),
+        domain=ir.FunCall(fun=ir.SymRef(id="cartesian_domain"), args=[]),
         stencil=ir.Lambda(
             params=[ir.Sym(id="y")],
             expr=ir.FunCall(fun=ir.SymRef(id="deref"), args=[ir.SymRef(id="y")]),
@@ -27,7 +27,7 @@ def test_simple():
 
 def test_shadowing():
     testee = ir.StencilClosure(
-        domain=ir.SymRef(id="d"),
+        domain=ir.FunCall(fun=ir.SymRef(id="cartesian_domain"), args=[]),
         stencil=ir.Lambda(
             params=[ir.Sym(id="x"), ir.Sym(id="y"), ir.Sym(id="z")],
             expr=ir.FunCall(
@@ -42,7 +42,7 @@ def test_shadowing():
         inputs=[ir.SymRef(id="foo"), ir.SymRef(id="bar"), ir.SymRef(id="baz")],
     )
     expected = ir.StencilClosure(
-        domain=ir.SymRef(id="d"),
+        domain=ir.FunCall(fun=ir.SymRef(id="cartesian_domain"), args=[]),
         stencil=ir.Lambda(
             params=[ir.Sym(id="y")],
             expr=ir.FunCall(
diff --git a/tests/functional_tests/iterator_tests/test_trace_shifts.py b/tests/functional_tests/iterator_tests/test_trace_shifts.py
index 6e68e086..7c8d3ec2 100644
--- a/tests/functional_tests/iterator_tests/test_trace_shifts.py
+++ b/tests/functional_tests/iterator_tests/test_trace_shifts.py
@@ -7,7 +7,7 @@ def test_trivial():
         stencil=ir.SymRef(id="deref"),
         inputs=[ir.SymRef(id="inp")],
         output=ir.SymRef(id="out"),
-        domain=ir.SymRef(id="dom"),
+        domain=ir.FunCall(fun=ir.SymRef(id="cartesian_domain"), args=[]),
     )
     expected = {"inp": [()]}
 
@@ -35,7 +35,7 @@ def test_shift():
         ),
         inputs=[ir.SymRef(id="inp")],
         output=ir.SymRef(id="out"),
-        domain=ir.SymRef(id="dom"),
+        domain=ir.FunCall(fun=ir.SymRef(id="cartesian_domain"), args=[]),
     )
     expected = {"inp": [(ir.OffsetLiteral(value="I"), ir.OffsetLiteral(value=1))]}
 
@@ -68,7 +68,7 @@ def test_lift():
         ),
         inputs=[ir.SymRef(id="inp")],
         output=ir.SymRef(id="out"),
-        domain=ir.SymRef(id="dom"),
+        domain=ir.FunCall(fun=ir.SymRef(id="cartesian_domain"), args=[]),
     )
     expected = {"inp": [(ir.OffsetLiteral(value="I"), ir.OffsetLiteral(value=1))]}
 
@@ -85,7 +85,7 @@ def test_reduce():
         ),
         inputs=[ir.SymRef(id="inp")],
         output=ir.SymRef(id="out"),
-        domain=ir.SymRef(id="dom"),
+        domain=ir.FunCall(fun=ir.SymRef(id="cartesian_domain"), args=[]),
     )
     expected = {"inp": [(ALL_NEIGHBORS,)]}
 
diff --git a/tests/functional_tests/iterator_tests/test_unroll_reduce.py b/tests/functional_tests/iterator_tests/test_unroll_reduce.py
index 10a57c6a..18f135c1 100644
--- a/tests/functional_tests/iterator_tests/test_unroll_reduce.py
+++ b/tests/functional_tests/iterator_tests/test_unroll_reduce.py
@@ -7,6 +7,11 @@ from functional.iterator import ir
 from functional.iterator.transforms.unroll_reduce import UnrollReduce
 
 
+@pytest.fixture(params=[True, False])
+def has_skip_values(request):
+    return request.param
+
+
 @pytest.fixture
 def basic_reduction():
     UIDs.reset_sequence()
@@ -91,7 +96,7 @@ def reduction_with_irrelevant_full_shift():
     )
 
 
-def _expected(red, dim, max_neighbors, has_skip_values):
+def _expected(red, dim, max_neighbors, has_skip_values, shifted_arg=0):
     acc = ir.SymRef(id="_acc_1")
     offset = ir.SymRef(id="_i_2")
     step = ir.SymRef(id="_step_3")
@@ -118,7 +123,7 @@ def _expected(red, dim, max_neighbors, has_skip_values):
             args=[
                 ir.FunCall(
                     fun=ir.FunCall(fun=ir.SymRef(id="shift"), args=[offset]),
-                    args=[red.args[0]],
+                    args=[red.args[shifted_arg]],
                 )
             ],
         )
@@ -132,26 +137,18 @@ def _expected(red, dim, max_neighbors, has_skip_values):
     return ir.FunCall(fun=ir.Lambda(params=[ir.Sym(id=step.id)], expr=step_app), args=[step_fun])
 
 
-def test_no_skip_values(basic_reduction):
-    expected = _expected(basic_reduction, "dim", 3, False)
-
-    offset_provider = {"dim": SimpleNamespace(max_neighbors=3, has_skip_values=False)}
-    actual = UnrollReduce.apply(basic_reduction, offset_provider=offset_provider)
-    assert actual == expected
-
-
-def test_skip_values(basic_reduction):
-    expected = _expected(basic_reduction, "dim", 3, True)
+def test_basic(basic_reduction, has_skip_values):
+    expected = _expected(basic_reduction, "dim", 3, has_skip_values)
 
-    offset_provider = {"dim": SimpleNamespace(max_neighbors=3, has_skip_values=True)}
+    offset_provider = {"dim": SimpleNamespace(max_neighbors=3, has_skip_values=has_skip_values)}
     actual = UnrollReduce.apply(basic_reduction, offset_provider=offset_provider)
     assert actual == expected
 
 
-def test_reduction_with_shift_on_second_arg(reduction_with_shift_on_second_arg):
-    expected = _expected(reduction_with_shift_on_second_arg, "dim", 3, False)
+def test_reduction_with_shift_on_second_arg(reduction_with_shift_on_second_arg, has_skip_values):
+    expected = _expected(reduction_with_shift_on_second_arg, "dim", 1, has_skip_values, 1)
 
-    offset_provider = {"dim": SimpleNamespace(max_neighbors=3, has_skip_values=False)}
+    offset_provider = {"dim": SimpleNamespace(max_neighbors=1, has_skip_values=has_skip_values)}
     actual = UnrollReduce.apply(reduction_with_shift_on_second_arg, offset_provider=offset_provider)
     assert actual == expected