Repository: GridTools/gt4py
Pull request #983 (2022-10-17 14:26:20)

Fix typing: iterator.embedded
## Description

Best effort to fix typing annotations in `iterator.embedded`.

## Requirements

Before submitting this PR, please make sure:

- [ ] You have run the code checks, tests and documentation build successfully
- [X] All fixes and all new functionality are tested and documentation is up to date
- [X] You looked at the [review checklist](https://github.com/GridTools/gt4py/blob/master/REVIEW_CHECKLIST.txt)

Additionally, if this PR contains code authored by new contributors:

- [X] All the authors are covered by a valid contributor assignment agreement,
signed by the employer if needed, provided to ETH Zurich
- [X] The names of all the new contributors have been added to an updated
version of the AUTHORS.rst file included in the PR
 





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..4bddd05f 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
 
@@ -236,13 +233,25 @@ ignore_errors = True
 [mypy-functional.type_system.type_translation]
 ignore_errors = True
 
-[mypy-functional.iterator.embedded]
+[mypy-functional.iterator.runtime]
+ignore_errors = True
+
+[mypy-functional.iterator.transforms.global_tmps]
 ignore_errors = True
 
-[mypy-functional.iterator.runtime]
+[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.*]
+[mypy-functional.iterator.transforms.inline_lambdas]
 ignore_errors = True
 
 #-- pytest --
diff --git a/src/functional/common.py b/src/functional/common.py
index 77e42c1a..331eab31 100644
--- a/src/functional/common.py
+++ b/src/functional/common.py
@@ -17,7 +17,9 @@ from __future__ import annotations
 import enum
 from collections.abc import Sequence
 from dataclasses import dataclass
-from typing import Any, Generic, Protocol, TypeVar, runtime_checkable
+from typing import Any, Generic, Optional, Protocol, TypeVar, runtime_checkable
+
+import numpy as np
 
 from eve.type_definitions import StrEnum
 
@@ -74,6 +76,12 @@ class Backend:
 class Connectivity(Protocol):
     max_neighbors: int
     has_skip_values: bool
+    origin_axis: Dimension
+
+    def mapped_index(
+        self, cur_index: int | np.integer, neigh_index: int | np.integer
+    ) -> Optional[int | np.integer]:
+        """Return neighbor index."""
 
 
 @enum.unique
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/embedded.py b/src/functional/iterator/embedded.py
index 601e4c3f..fd84c6d2 100644
--- a/src/functional/iterator/embedded.py
+++ b/src/functional/iterator/embedded.py
@@ -2,13 +2,12 @@
 
 from __future__ import annotations
 
+import abc
 import copy
+import dataclasses
 import itertools
 import math
 import numbers
-from abc import abstractmethod
-from dataclasses import dataclass
-from types import NoneType
 from typing import (
     Any,
     Callable,
@@ -22,6 +21,7 @@ from typing import (
     TypeAlias,
     TypeGuard,
     Union,
+    cast,
     runtime_checkable,
 )
 
@@ -29,11 +29,8 @@ import numpy as np
 import numpy.typing as npt
 
 from eve import extended_typing as xtyping
-from functional import iterator
-from functional.common import Connectivity, Dimension, DimensionKind
-from functional.iterator import builtins
-from functional.iterator.runtime import CartesianDomain, Offset, UnstructuredDomain
-from functional.iterator.utils import tupelize
+from functional import common
+from functional.iterator import builtins, runtime, utils
 
 
 EMBEDDED = "embedded"
@@ -46,11 +43,12 @@ IntIndex: TypeAlias = int | np.integer
 FieldIndex: TypeAlias = slice | IntIndex
 FieldIndexOrIndices: TypeAlias = FieldIndex | tuple[FieldIndex, ...]
 
+
 FieldAxis: TypeAlias = (
-    Dimension | Offset
+    common.Dimension | runtime.Offset
 )  # TODO Offset should be removed, is sometimes used for sparse dimensions
-TupleAxis: TypeAlias = NoneType
-Axis: TypeAlias = FieldAxis | TupleAxis
+TupleAxis: TypeAlias = type[None]
+Axis: TypeAlias = Union[FieldAxis, TupleAxis]
 Scalar: TypeAlias = SupportsInt | SupportsFloat | np.int32 | np.int64 | np.float32 | np.float64
 
 
@@ -62,24 +60,27 @@ class NeighborTableOffsetProvider:
     def __init__(
         self,
         tbl: npt.NDArray,  # TODO(havogt): define neighbor table concept
-        origin_axis: Dimension,
-        neighbor_axis: Dimension,
+        origin_axis: common.Dimension,
+        neighbor_axis: common.Dimension,
         max_neighbors: int,
         has_skip_values=True,
     ) -> None:
-        self.tbl = tbl
+        self._tbl = tbl
         self.origin_axis = origin_axis
         self.neighbor_axis = neighbor_axis
         assert not hasattr(tbl, "shape") or tbl.shape[1] == max_neighbors
         self.max_neighbors = max_neighbors
         self.has_skip_values = has_skip_values
 
+    def mapped_index(self, primary: IntIndex, neighbor_idx: IntIndex) -> IntIndex:
+        return self._tbl[(primary, neighbor_idx)]
+
 
 class StridedNeighborOffsetProvider:
     def __init__(
         self,
-        origin_axis: Dimension,
-        neighbor_axis: Dimension,
+        origin_axis: common.Dimension,
+        neighbor_axis: common.Dimension,
         max_neighbors: int,
         has_skip_values=True,
     ) -> None:
@@ -88,23 +89,14 @@ class StridedNeighborOffsetProvider:
         self.max_neighbors = max_neighbors
         self.has_skip_values = has_skip_values
 
-    @property
-    def tbl(self):  # TODO(havogt): define Connectivity concept properly
-        class Impl:
-            def __init__(self, stride: int) -> None:
-                self.stride = stride
-
-            def __getitem__(self, indices: tuple[int, int]) -> int:
-                primary, neighbor_idx = indices
-                return primary * self.stride + neighbor_idx
-
-        return Impl(stride=self.max_neighbors)
+    def mapped_index(self, primary: IntIndex, neighbor_idx: IntIndex) -> IntIndex:
+        return primary * self.max_neighbors + neighbor_idx
 
 
 # Offsets
 OffsetPart: TypeAlias = Tag | IntIndex
 CompleteOffset: TypeAlias = tuple[Tag, IntIndex]
-OffsetProviderElem: TypeAlias = Dimension | Connectivity
+OffsetProviderElem: TypeAlias = common.Dimension | common.Connectivity
 OffsetProvider: TypeAlias = dict[Tag, OffsetProviderElem]
 
 # Positions
@@ -150,11 +142,11 @@ class LocatedField(Protocol):
     """A field with named dimensions providing read access."""
 
     @property
-    @abstractmethod
-    def axes(self) -> tuple[Dimension, ...]:
+    @abc.abstractmethod
+    def axes(self) -> tuple[common.Dimension, ...]:
         ...
 
-    @abstractmethod
+    @abc.abstractmethod
     def __getitem__(self, indices: FieldIndexOrIndices) -> Any:
         ...
 
@@ -162,11 +154,11 @@ class LocatedField(Protocol):
 class MutableLocatedField(LocatedField, Protocol):
     """A LocatedField with write access."""
 
-    @abstractmethod
+    @abc.abstractmethod
     def __setitem__(self, indices: FieldIndexOrIndices, value: Any) -> None:
         ...
 
-    @abstractmethod
+    @abc.abstractmethod
     def __array__(self) -> np.ndarray:
         ...
 
@@ -305,7 +297,7 @@ def lift(stencil):
     def impl(*args):
         class _WrappedIterator:
             def __init__(
-                self, stencil, args, *, offsets: list[OffsetPart] = None, elem=None
+                self, stencil, args, *, offsets: Optional[list[OffsetPart]] = None, elem=None
             ) -> None:
                 assert not offsets or all(isinstance(o, (int, str)) for o in offsets)
                 self.stencil = stencil
@@ -385,24 +377,26 @@ def reduce(fun, init):
     return sten
 
 
-NamedRange: TypeAlias = tuple[Tag | Dimension, range]
+NamedRange: TypeAlias = tuple[Tag | common.Dimension, range]
 
 
 @builtins.cartesian_domain.register(EMBEDDED)
-def cartesian_domain(*args: NamedRange) -> CartesianDomain:
-    return CartesianDomain(args)
+def cartesian_domain(*args: NamedRange) -> runtime.CartesianDomain:
+    return runtime.CartesianDomain(args)
 
 
 @builtins.unstructured_domain.register(EMBEDDED)
-def unstructured_domain(*args: NamedRange) -> UnstructuredDomain:
-    return UnstructuredDomain(args)
+def unstructured_domain(*args: NamedRange) -> runtime.UnstructuredDomain:
+    return runtime.UnstructuredDomain(args)
 
 
-Domain: TypeAlias = CartesianDomain | UnstructuredDomain | dict[str | Dimension, range]
+Domain: TypeAlias = (
+    runtime.CartesianDomain | runtime.UnstructuredDomain | dict[str | common.Dimension, range]
+)
 
 
 @builtins.named_range.register(EMBEDDED)
-def named_range(tag: Tag | Dimension, start: int, end: int) -> NamedRange:
+def named_range(tag: Tag | common.Dimension, start: int, end: int) -> NamedRange:
     # TODO revisit this pattern after the discussion of 0d-field vs scalar
     if isinstance(start, ConstantField):
         start = start.value
@@ -474,15 +468,25 @@ def not_eq(first, second):
 CompositeOfScalarOrField: TypeAlias = Scalar | LocatedField | tuple["CompositeOfScalarOrField", ...]
 
 
+def is_dtype_like(t: Any) -> TypeGuard[npt.DTypeLike]:
+    return issubclass(t, (np.generic, int, float))
+
+
+def infer_dtype_like_type(t: Any) -> npt.DTypeLike:
+    res = xtyping.infer_type(t)
+    assert is_dtype_like(res), res
+    return res
+
+
 def promote_scalars(val: CompositeOfScalarOrField):
     """Given a scalar, field or composite thereof promote all (contained) scalars to fields."""
     if isinstance(val, tuple):
         return tuple(promote_scalars(el) for el in val)
-    val_type = xtyping.infer_type(val)
+    elif isinstance(val, LocatedField):
+        return val
+    val_type = infer_dtype_like_type(val)
     if np.issubdtype(val_type, np.number):
         return constant_field(val)
-    elif np.issubdtype(val_type, LocatedField):
-        return val
     else:
         raise ValueError(
             f"Expected a `Field` or a number (`float`, `np.int64`, ...), but got {val_type}."
@@ -492,6 +496,7 @@ def promote_scalars(val: CompositeOfScalarOrField):
 for math_builtin_name in builtins.MATH_BUILTINS:
     python_builtins = {"int": int, "float": float, "bool": bool, "str": str}
     decorator = getattr(builtins, math_builtin_name).register(EMBEDDED)
+    impl: Callable
     if math_builtin_name == "gamma":
         # numpy has no gamma function
         impl = np.vectorize(math.gamma)
@@ -511,8 +516,10 @@ def _lookup_offset_provider(offset_provider: OffsetProvider, tag: Tag) -> Offset
     return offset_provider[tag]
 
 
-def _get_connectivity(offset_provider: OffsetProvider, tag: Tag) -> Connectivity:
-    if not isinstance(connectivity := _lookup_offset_provider(offset_provider, tag), Connectivity):
+def _get_connectivity(offset_provider: OffsetProvider, tag: Tag) -> common.Connectivity:
+    if not isinstance(
+        connectivity := _lookup_offset_provider(offset_provider, tag), common.Connectivity
+    ):
         raise RuntimeError(f"Expected a `Connectivity` for `{tag}`")
     return connectivity
 
@@ -521,7 +528,7 @@ def _named_range(axis: str, range_: Iterable[int]) -> Iterable[CompleteOffset]:
     return ((axis, i) for i in range_)
 
 
-def _domain_iterator(domain: dict[str, range]) -> Iterable[Position]:
+def _domain_iterator(domain: dict[Tag, range]) -> Iterable[Position]:
     return (
         dict(elem)
         for elem in itertools.product(*(_named_range(axis, rang) for axis, rang in domain.items()))
@@ -537,16 +544,20 @@ def execute_shift(
         assert isinstance(current_entry, list)
         new_entry = list(current_entry)
         assert None in new_entry
+        assert isinstance(
+            index, int
+        )  # narrowing to `int` as it's an element of `SparsePositionEntry`
         for i, p in reversed(list(enumerate(new_entry))):
             # first shift applies to the last sparse dimensions of that axis type
             if p is None:
                 new_entry[i] = index
                 break
-        return pos | {tag: new_entry}
+        # the assertions above confirm pos is incomplete casting here to avoid duplicating work in a type guard
+        return cast(IncompletePosition, pos) | {tag: new_entry}
 
     assert tag in offset_provider
     offset_implementation = offset_provider[tag]
-    if isinstance(offset_implementation, Dimension):
+    if isinstance(offset_implementation, common.Dimension):
         new_pos = copy.copy(pos)
         if is_int_index(value := new_pos[offset_implementation.value]):
             new_pos[offset_implementation.value] = value + index
@@ -554,19 +565,22 @@ def execute_shift(
             raise AssertionError()
         return new_pos
     else:
-        assert isinstance(offset_implementation, Connectivity)
+        assert isinstance(offset_implementation, common.Connectivity)
         assert offset_implementation.origin_axis.value in pos
         new_pos = pos.copy()
         new_pos.pop(offset_implementation.origin_axis.value)
-        if offset_implementation.tbl[pos[offset_implementation.origin_axis.value], index] in [
+        cur_index = pos[offset_implementation.origin_axis.value]
+        assert is_int_index(cur_index)
+        if offset_implementation.mapped_index(cur_index, index) in [
             None,
             -1,
         ]:
             return None
         else:
-            new_pos[offset_implementation.neighbor_axis.value] = int(
-                offset_implementation.tbl[pos[offset_implementation.origin_axis.value], index]
-            )
+            new_index = offset_implementation.mapped_index(cur_index, index)
+            assert new_index is not None
+            new_pos[offset_implementation.neighbor_axis.value] = int(new_index)
+
         return new_pos
 
     raise AssertionError("Unknown object in `offset_provider`")
@@ -681,7 +695,7 @@ def _is_concrete_position(pos: Position) -> TypeGuard[ConcretePosition]:
 
 def _get_axes(
     field_or_tuple: LocatedField | tuple,
-) -> Sequence[Dimension]:  # arbitrary nesting of tuples of LocatedField
+) -> Sequence[common.Dimension]:  # arbitrary nesting of tuples of LocatedField
     return (
         _get_axes(field_or_tuple[0]) if isinstance(field_or_tuple, tuple) else field_or_tuple.axes
     )
@@ -689,10 +703,10 @@ def _get_axes(
 
 def _make_tuple(
     field_or_tuple: LocatedField | tuple,  # arbitrary nesting of tuples of LocatedField
-    indices: int | slice | tuple[int | slice, ...],
+    indices: FieldIndexOrIndices,
     *,
     as_column: bool = False,
-) -> tuple:  # arbitrary nesting of tuples of field values or `Column`s
+) -> tuple | Column:  # arbitrary nesting of tuples of field values or `Column`s
     if isinstance(field_or_tuple, tuple):
         return tuple(_make_tuple(f, indices, as_column=as_column) for f in field_or_tuple)
     else:
@@ -706,22 +720,13 @@ def _make_tuple(
             return data
 
 
-# TODO(havogt) frozen dataclass
+@dataclasses.dataclass(frozen=True)
 class MDIterator:
-    def __init__(
-        self,
-        field: LocatedField,
-        pos: MaybePosition,
-        *,
-        incomplete_offsets: Sequence[Tag] = None,
-        offset_provider: OffsetProvider,
-        column_axis: Tag = None,
-    ) -> None:
-        self.field = field
-        self.pos = pos
-        self.incomplete_offsets = incomplete_offsets or []
-        self.offset_provider = offset_provider
-        self.column_axis = column_axis
+    field: LocatedField
+    pos: MaybePosition
+    incomplete_offsets: Sequence[Tag] = dataclasses.field(default_factory=list, kw_only=True)
+    offset_provider: OffsetProvider = dataclasses.field(kw_only=True)
+    column_axis: Optional[Tag] = dataclasses.field(default=None, kw_only=True)
 
     def shift(self, *offsets: OffsetPart) -> MDIterator:
         complete_offsets, open_offsets = group_offsets(*self.incomplete_offsets, *offsets)
@@ -780,10 +785,10 @@ def make_in_iterator(
     axes = _get_axes(inp)
     sparse_dimensions: list[Tag] = []
     for axis in axes:
-        if isinstance(axis, Offset):
+        if isinstance(axis, runtime.Offset):
             assert isinstance(axis.value, str)
             sparse_dimensions.append(axis.value)
-        elif isinstance(axis, Dimension) and axis.kind == DimensionKind.LOCAL:
+        elif isinstance(axis, common.Dimension) and axis.kind == common.DimensionKind.LOCAL:
             # we just use the name of the axis to match the offset literal for now
             sparse_dimensions.append(axis.value)
 
@@ -810,13 +815,13 @@ class LocatedFieldImpl(MutableLocatedField):
     """A Field with named dimensions/axes."""
 
     @property
-    def axes(self) -> tuple[Dimension, ...]:
+    def axes(self) -> tuple[common.Dimension, ...]:
         return self._axes
 
     def __init__(
         self,
         getter: Callable[[FieldIndexOrIndices], Any],
-        axes: tuple[Dimension, ...],
+        axes: tuple[common.Dimension, ...],
         dtype,
         *,
         setter: Callable[[FieldIndexOrIndices, Any], None],
@@ -829,7 +834,7 @@ class LocatedFieldImpl(MutableLocatedField):
         self.dtype = dtype
 
     def __getitem__(self, indices: FieldIndexOrIndices) -> Any:
-        indices = tupelize(indices)
+        indices = utils.tupelize(indices)
         return self.getter(indices)
 
     def __setitem__(self, indices: FieldIndexOrIndices, value: Any):
@@ -845,14 +850,14 @@ class LocatedFieldImpl(MutableLocatedField):
         return self.array().shape
 
 
-def _is_tuple_axis(axis: Axis) -> TypeGuard[TupleAxis]:
-    return isinstance(axis, TupleAxis)
-
-
 def _is_field_axis(axis: Axis) -> TypeGuard[FieldAxis]:
     return isinstance(axis, FieldAxis)  # type: ignore[misc,arg-type] # see https://github.com/python/mypy/issues/11673
 
 
+def _is_tuple_axis(axis: Axis) -> TypeGuard[TupleAxis]:
+    return axis is None
+
+
 def _is_sparse_position_entry(
     pos: FieldIndex | SparsePositionEntry,
 ) -> TypeGuard[SparsePositionEntry]:
@@ -870,6 +875,7 @@ def get_ordered_indices(
         else:
             assert _is_field_axis(axis)
             assert axis.value in pos
+            assert isinstance(axis.value, str)
             elem = pos[axis.value]
             if _is_sparse_position_entry(elem):
                 sparse_position_tracker.setdefault(axis.value, 0)
@@ -907,7 +913,7 @@ def _tupsum(a, b):
 
 
 def np_as_located_field(
-    *axes: Dimension, origin: Optional[dict[Dimension, int]] = None
+    *axes: common.Dimension, origin: Optional[dict[common.Dimension, int]] = None
 ) -> Callable[[np.ndarray], LocatedFieldImpl]:
     def _maker(a: np.ndarray) -> LocatedFieldImpl:
         if a.ndim != len(axes):
@@ -919,7 +925,7 @@ def np_as_located_field(
             offsets = None
 
         def setter(indices, value):
-            indices = tupelize(indices)
+            indices = utils.tupelize(indices)
             a[_tupsum(indices, offsets) if offsets else indices] = value
 
         def getter(indices):
@@ -931,20 +937,23 @@ def np_as_located_field(
 
 
 class IndexField(LocatedField):
-    def __init__(self, axis: Dimension, dtype: npt.DTypeLike) -> None:
+    def __init__(self, axis: common.Dimension, dtype: npt.DTypeLike) -> None:
         self.axis = axis
         self.dtype = np.dtype(dtype)
 
     def __getitem__(self, index: FieldIndexOrIndices) -> Any:
-        assert isinstance(index, int) or (isinstance(index, tuple) and len(index) == 1)
-        return self.dtype.type(index if isinstance(index, int) else index[0])
+        if isinstance(index, int):
+            return self.dtype.type(index)
+        else:
+            assert isinstance(index, tuple) and len(index) == 1 and isinstance(index[0], int)
+            return self.dtype.type(index[0])
 
     @property
-    def axes(self) -> tuple[Dimension]:
+    def axes(self) -> tuple[common.Dimension]:
         return (self.axis,)
 
 
-def index_field(axis: Dimension, dtype: npt.DTypeLike = int) -> LocatedField:
+def index_field(axis: common.Dimension, dtype: npt.DTypeLike = int) -> LocatedField:
     return IndexField(axis, dtype)
 
 
@@ -961,33 +970,31 @@ class ConstantField(LocatedField):
         return ()
 
 
-def constant_field(value: Any, dtype: npt.DTypeLike = None) -> LocatedField:
+def constant_field(value: Any, dtype: Optional[npt.DTypeLike] = None) -> LocatedField:
     if dtype is None:
-        dtype = xtyping.infer_type(value)
+        dtype = infer_dtype_like_type(value)
     return ConstantField(value, dtype)
 
 
 @builtins.shift.register(EMBEDDED)
-def shift(*offsets: Union[Offset, int]) -> Callable[[ItIterator], ItIterator]:
+def shift(*offsets: Union[runtime.Offset, int]) -> Callable[[ItIterator], ItIterator]:
     def impl(it: ItIterator) -> ItIterator:
-        return it.shift(*list(o.value if isinstance(o, Offset) else o for o in offsets))
+        return it.shift(*list(o.value if isinstance(o, runtime.Offset) else o for o in offsets))
 
     return impl
 
 
-@dataclass
+@dataclasses.dataclass(frozen=True)
 class ColumnDescriptor:
     axis: str
     col_range: range  # TODO(havogt) introduce range type that doesn't have step
 
 
+@dataclasses.dataclass(frozen=True)
 class ScanArgIterator:
-    def __init__(
-        self, wrapped_iter: ItIterator, k_pos: int, *, offsets: Sequence[OffsetPart] = None
-    ) -> None:
-        self.wrapped_iter = wrapped_iter
-        self.offsets = offsets or []
-        self.k_pos = k_pos
+    wrapped_iter: ItIterator
+    k_pos: int
+    offsets: Sequence[OffsetPart] = dataclasses.field(default_factory=list, kw_only=True)
 
     def deref(self) -> Any:
         if not self.can_deref():
@@ -1132,12 +1139,12 @@ def scan(scan_pass, is_forward: bool, init):
 
 
 def _dimension_to_tag(domain: Domain) -> dict[Tag, range]:
-    return {k.value if isinstance(k, Dimension) else k: v for k, v in domain.items()}
+    return {k.value if isinstance(k, common.Dimension) else k: v for k, v in domain.items()}
 
 
 def _validate_domain(domain: Domain, offset_provider: OffsetProvider) -> None:
-    if isinstance(domain, CartesianDomain):
-        if any(isinstance(o, Connectivity) for o in offset_provider.values()):
+    if isinstance(domain, runtime.CartesianDomain):
+        if any(isinstance(o, common.Connectivity) for o in offset_provider.values()):
             raise RuntimeError(
                 "Got a `CartesianDomain`, but found a `Connectivity` in `offset_provider`, expected `UnstructuredDomain`."
             )
@@ -1147,7 +1154,7 @@ def fendef_embedded(fun: Callable[..., None], *args: Any, **kwargs: Any):
     if "offset_provider" not in kwargs:
         raise RuntimeError("offset_provider not provided")
 
-    @iterator.runtime.closure.register(EMBEDDED)
+    @runtime.closure.register(EMBEDDED)
     def closure(
         domain_: Domain,
         sten: Callable[..., Any],
@@ -1200,4 +1207,4 @@ def fendef_embedded(fun: Callable[..., None], *args: Any, **kwargs: Any):
     fun(*args)
 
 
-iterator.runtime.fendef_embedded = fendef_embedded
+runtime.fendef_embedded = fendef_embedded
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..d401fbeb 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,
@@ -362,11 +364,12 @@ def _max_domain_sizes_by_location_type(offset_provider: Mapping[str, Any]) -> di
             assert provider.origin_axis.kind == DimensionKind.HORIZONTAL
             assert provider.neighbor_axis.kind == DimensionKind.HORIZONTAL
             sizes[provider.origin_axis.value] = max(
-                sizes.get(provider.origin_axis.value, 0), provider.tbl.shape[0]
+                sizes.get(provider.origin_axis.value, 0),
+                provider._tbl.shape[0],  # TODO properly expose the size
             )
             sizes[provider.neighbor_axis.value] = max(
                 sizes.get(provider.neighbor_axis.value, 0),
-                provider.tbl.max(),
+                provider._tbl.max(),
             )
     return sizes
 
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/cpp_backend_tests/fvm_nabla.py b/tests/functional_tests/cpp_backend_tests/fvm_nabla.py
index 0d2e7a71..4d830ff1 100644
--- a/tests/functional_tests/cpp_backend_tests/fvm_nabla.py
+++ b/tests/functional_tests/cpp_backend_tests/fvm_nabla.py
@@ -1,4 +1,5 @@
 import sys
+from dataclasses import dataclass
 from types import SimpleNamespace
 
 from functional.common import Dimension, DimensionKind
@@ -60,6 +61,16 @@ def nabla_fencil(n_vertices, n_levels, out, pp, S_M, sign, vol):
     )
 
 
+@dataclass
+class DummyConnectivity:
+    max_neighbors: int
+    has_skip_values: int
+    origin_axis: Dimension = Dimension("dummy")
+
+    def mapped_index(_, __) -> int:
+        return 0
+
+
 if __name__ == "__main__":
     if len(sys.argv) != 2:
         raise RuntimeError(f"Usage: {sys.argv[0]} <output_file>")
@@ -68,8 +79,8 @@ if __name__ == "__main__":
     # prog = trace(zavgS_fencil, [None] * 4) # TODO allow generating of 2 fencils
     prog = trace(nabla_fencil, [None] * 7)
     offset_provider = {
-        "V2E": SimpleNamespace(max_neighbors=6, has_skip_values=True),
-        "E2V": SimpleNamespace(max_neighbors=2, has_skip_values=False),
+        "V2E": DummyConnectivity(max_neighbors=6, has_skip_values=True),
+        "E2V": DummyConnectivity(max_neighbors=2, has_skip_values=False),
     }
     generated_code = generate(prog, offset_provider=offset_provider)
 
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_domain.py b/tests/functional_tests/iterator_tests/test_domain.py
index a184c587..abad9de3 100644
--- a/tests/functional_tests/iterator_tests/test_domain.py
+++ b/tests/functional_tests/iterator_tests/test_domain.py
@@ -1,4 +1,4 @@
-from types import SimpleNamespace
+from dataclasses import dataclass
 
 import numpy as np
 import pytest
@@ -14,7 +14,17 @@ def foo(inp):
     return deref(inp)
 
 
-connectivity = SimpleNamespace(max_neighbors=0, has_skip_values=True)
+@dataclass
+class DummyConnectivity:
+    max_neighbors: int
+    has_skip_values: int
+    origin_axis: Dimension = Dimension("dummy")
+
+    def mapped_index(_, __) -> int:
+        return 0
+
+
+connectivity = DummyConnectivity(max_neighbors=0, has_skip_values=True)
 
 
 def test_deduce_domain():
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