Repository: GridTools/gt4py
Pull request #1117 (2023-01-05 10:06:03)

Fix mypy gtfn_backend
Un-ignore `gtfn_backend` and fix all the reported mypy errors



Patch:
diff --git a/setup.cfg b/setup.cfg
index fce9e04c..a0070073 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -221,8 +221,6 @@ disallow_incomplete_defs = False
 
 ## TODO: the following ignore_errors are only temporary, while
 # we work in fixing the wrong typings
-[mypy-functional.program_processors.codegens/gtfn/gtfn_backend]
-ignore_errors = True
 
 [mypy-functional.ffront.decorator]
 ignore_errors = True
@@ -233,13 +231,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.runtime]
+[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.*]
+[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/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/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 3080ac20..d401fbeb 100644
--- a/src/functional/iterator/transforms/global_tmps.py
+++ b/src/functional/iterator/transforms/global_tmps.py
@@ -364,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/gtfn_backend.py b/src/functional/program_processors/codegens/gtfn/gtfn_backend.py
index a5b2e092..ecdf146b 100644
--- a/src/functional/program_processors/codegens/gtfn/gtfn_backend.py
+++ b/src/functional/program_processors/codegens/gtfn/gtfn_backend.py
@@ -25,6 +25,7 @@ from functional.program_processors.codegens.gtfn.itir_to_gtfn_ir import GTFN_low
 def generate(program: itir.FencilDefinition, **kwargs: Any) -> str:
     transformed = program
     offset_provider = kwargs.get("offset_provider")
+    assert isinstance(offset_provider, dict)
     transformed = apply_common_transforms(
         program,
         lift_mode=kwargs.get("lift_mode"),
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 1fab4e47..478cd223 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
@@ -199,7 +199,7 @@ class GTFN_lowering(eve.NodeTranslator, eve.VisitorWithSymbolTableTrait):
     _unary_op_map: ClassVar[dict[str, str]] = {"not_": "!"}
 
     offset_provider: dict
-    column_axis: common.Dimension
+    column_axis: Optional[common.Dimension]
     grid_type: common.GridType
 
     # we use one UID generator per instance such that the generated ids are
@@ -212,7 +212,7 @@ class GTFN_lowering(eve.NodeTranslator, eve.VisitorWithSymbolTableTrait):
         node: itir.FencilDefinition | FencilWithTemporaries,
         *,
         offset_provider: dict,
-        column_axis: common.Dimension,
+        column_axis: Optional[common.Dimension],
     ):
         if isinstance(node, FencilWithTemporaries):
             fencil_definition = node.fencil
@@ -457,11 +457,13 @@ class GTFN_lowering(eve.NodeTranslator, eve.VisitorWithSymbolTableTrait):
                 inputs=[Literal(value=str(i + 1), type="int") for i, _ in enumerate(node.inputs)],
                 init=self.visit(node.stencil.args[2], **kwargs),
             )
+            column_axis = self.column_axis
+            assert isinstance(column_axis, common.Dimension)
             return ScanExecution(
                 backend=backend,
                 scans=[scan],
                 args=[self.visit(node.output, **kwargs)] + self.visit(node.inputs),
-                axis=SymRef(id=self.column_axis.value),
+                axis=SymRef(id=column_axis.value),
             )
         return StencilExecution(
             stencil=self.visit(


Test patch:
diff --git a/tests/functional_tests/cpp_backend_tests/copy_stencil.py b/tests/functional_tests/cpp_backend_tests/copy_stencil.py
index 922d9876..4535df70 100644
--- a/tests/functional_tests/cpp_backend_tests/copy_stencil.py
+++ b/tests/functional_tests/cpp_backend_tests/copy_stencil.py
@@ -33,7 +33,7 @@ if __name__ == "__main__":
     output_file = sys.argv[1]
 
     prog = trace(copy_fencil, [None] * 5)
-    generated_code = generate(prog)
+    generated_code = generate(prog, offset_provider={})
 
     with open(output_file, "w+") as output:
         output.write(generated_code)
diff --git a/tests/functional_tests/cpp_backend_tests/copy_stencil_field_view.py b/tests/functional_tests/cpp_backend_tests/copy_stencil_field_view.py
index a8b13a94..ac1c5199 100644
--- a/tests/functional_tests/cpp_backend_tests/copy_stencil_field_view.py
+++ b/tests/functional_tests/cpp_backend_tests/copy_stencil_field_view.py
@@ -34,7 +34,7 @@ if __name__ == "__main__":
     output_file = sys.argv[1]
 
     prog = copy_program.itir
-    generated_code = generate(prog)
+    generated_code = generate(prog, offset_provider={})
 
     with open(output_file, "w+") as output:
         output.write(generated_code)
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/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_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