Repository: GridTools/gt4py
Pull request #1216 (2023-03-21 16:36:28)

Simplify ITIR shift collector (TraceShifts)
The `TraceShifts` pass had some bugs in particular related to literals and tuples. This PR simplifies the pass, fixes the bugs and adds some more tests.

tehrengruber: @havogt The CI was failing for unrelated reasons. Can you approve?

Patch:
diff --git a/src/gt4py/next/iterator/transforms/trace_shifts.py b/src/gt4py/next/iterator/transforms/trace_shifts.py
index 284fc2cc..5024c607 100644
--- a/src/gt4py/next/iterator/transforms/trace_shifts.py
+++ b/src/gt4py/next/iterator/transforms/trace_shifts.py
@@ -12,15 +12,20 @@
 #
 # SPDX-License-Identifier: GPL-3.0-or-later
 
+import types
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any, Final
+from typing import Any, Final, Union
 
 from gt4py.eve import NodeTranslator
 from gt4py.next.iterator import ir
 from gt4py.next.iterator.transforms.collect_shifts import ALL_NEIGHBORS
 
 
+VALUE_TOKEN = types.new_class("VALUE_TOKEN")
+TYPE_TOKEN = types.new_class("TYPE_TOKEN")
+
+
 @dataclass(frozen=True)
 class InputTracer:
     inp: str
@@ -37,77 +42,26 @@ class InputTracer:
         )
 
     def deref(self):
-        if self.lift_level:
-            return self
         self.register_deref(self.inp, self.offsets)
+        return VALUE_TOKEN
 
-    def lift(self):
-        return InputTracer(
-            inp=self.inp,
-            register_deref=self.register_deref,
-            offsets=self.offsets,
-            lift_level=self.lift_level + 1,
-        )
-
-    def unlift(self):
-        assert self.lift_level > 0
-        return InputTracer(
-            inp=self.inp,
-            register_deref=self.register_deref,
-            offsets=self.offsets,
-            lift_level=self.lift_level - 1,
-        )
-
-
-@dataclass(frozen=True)
-class CombinedTracer:
-    tracers: tuple[InputTracer, ...]
-
-    def shift(self, offsets):
-        return CombinedTracer(tracers=tuple(t.shift(offsets) for t in self.tracers))
-
-    def deref(self):
-        lift_levels = {t.lift_level for t in self.tracers}
-        assert len(lift_levels) == 1
-        if lift_levels.pop():
-            return self
-        for t in self.tracers:
-            t.deref()
-
-    def lift(self):
-        return CombinedTracer(tracers=tuple(t.lift() for t in self.tracers))
-
-    def unlift(self):
-        return CombinedTracer(tracers=tuple(t.unlift() for t in self.tracers))
-
-
-def _combine(*tracers):
-    input_tracers = []
-
-    def handle_tracer(tracer):
-        if isinstance(tracer, InputTracer):
-            input_tracers.append(tracer)
-        elif isinstance(tracer, CombinedTracer):
-            input_tracers.extend(tracer.tracers)
-        elif isinstance(tracer, tuple):
-            for t in tracer:
-                handle_tracer(t)
 
-    for tracer in tracers:
-        handle_tracer(tracer)
-
-    return CombinedTracer(tracers=tuple(input_tracers))
+def _combine(*values):
+    # `OffsetLiteral`s may occur in `list_get` calls
+    if not all(
+        val in [VALUE_TOKEN, TYPE_TOKEN] or isinstance(val, ir.OffsetLiteral) for val in values
+    ):
+        raise AssertionError("All arguments must be values or types.")
+    return VALUE_TOKEN
 
 
 # implementations of builtins
-
-
 def _deref(x):
     return x.deref()
 
 
 def _can_deref(x):
-    return
+    return VALUE_TOKEN
 
 
 def _shift(*offsets):
@@ -117,9 +71,23 @@ def _shift(*offsets):
     return apply
 
 
+@dataclass(frozen=True)
+class AppliedLift:
+    stencil: Callable
+    its: tuple[Union[InputTracer, "AppliedLift"]]
+
+    def shift(self, offsets):
+        return AppliedLift(self.stencil, tuple(_shift(it) for it in self.its))
+
+    def deref(self):
+        return self.stencil(*self.its)
+
+
 def _lift(f):
-    def apply(*args):
-        return f(*(arg.lift() for arg in args)).unlift()
+    def apply(*its):
+        if not all(isinstance(it, (InputTracer, AppliedLift)) for it in its):
+            raise AssertionError("All arguments must be iterators.")
+        return AppliedLift(f, its)
 
     return apply
 
@@ -128,8 +96,12 @@ def _reduce(f, init):
     return _combine
 
 
+def _map(f):
+    return _combine
+
+
 def _neighbors(o, x):
-    return x.shift((o, ALL_NEIGHBORS)).deref()
+    return _deref(_shift(o, ALL_NEIGHBORS)(x))
 
 
 def _scan(f, forward, init):
@@ -139,32 +111,27 @@ def _scan(f, forward, init):
     return apply
 
 
-def _make_tuple(*args):
-    return args
-
-
-def _tuple_get(idx, tup):
-    if tup is not None:
-        return tup[int(idx.value)]
-
-
 _START_CTX: Final = {
     "deref": _deref,
     "can_deref": _can_deref,
     "shift": _shift,
     "lift": _lift,
-    "reduce": _reduce,
     "scan": _scan,
-    "make_tuple": _make_tuple,
-    "tuple_get": _tuple_get,
+    "reduce": _reduce,
     "neighbors": _neighbors,
+    "map_": _map,
 }
 
 
 class TraceShifts(NodeTranslator):
+    def visit_Literal(self, node: ir.SymRef, *, ctx: dict[str, Any]) -> Any:
+        return VALUE_TOKEN
+
     def visit_SymRef(self, node: ir.SymRef, *, ctx: dict[str, Any]) -> Any:
         if node.id in ctx:
             return ctx[node.id]
+        elif node.id in ir.TYPEBUILTINS:
+            return TYPE_TOKEN
         return _combine
 
     def visit_FunCall(self, node: ir.FunCall, *, ctx: dict[str, Any]) -> Any:
@@ -182,10 +149,15 @@ class TraceShifts(NodeTranslator):
         self, node: ir.StencilClosure, *, shifts: dict[str, list[tuple[ir.OffsetLiteral, ...]]]
     ):
         def register_deref(inp: str, offsets: tuple[ir.OffsetLiteral, ...]):
-            shifts.setdefault(inp, []).append(offsets)
+            shifts[inp].append(offsets)
+
+        tracers = []
+        for inp in node.inputs:
+            shifts.setdefault(inp.id, [])
+            tracers.append(InputTracer(inp=inp.id, register_deref=register_deref))
 
-        tracers = [InputTracer(inp=inp.id, register_deref=register_deref) for inp in node.inputs]
-        self.visit(node.stencil, ctx=_START_CTX)(*tracers)
+        result = self.visit(node.stencil, ctx=_START_CTX)(*tracers)
+        assert result is VALUE_TOKEN
 
     @classmethod
     def apply(cls, node: ir.StencilClosure) -> dict[str, list[tuple[ir.OffsetLiteral, ...]]]:


Test patch:
diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py
index 82fe7d5f..072a5304 100644
--- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py
+++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py
@@ -26,8 +26,8 @@ from gt4py.next.ffront.fbuiltins import (
     Field,
     astype,
     broadcast,
+    float32,
     float64,
-    int32,
     int64,
     maximum,
     minimum,
@@ -61,8 +61,8 @@ def test_copy(fieldview_backend):
 def test_multicopy(fieldview_backend):
     a_I_float = np_as_located_field(IDim)(np.random.randn(size).astype("float64"))
     b_I_float = np_as_located_field(IDim)(np.random.randn(size).astype("float64"))
-    out_I_float = np_as_located_field(IDim)(np.zeros((size), dtype=float64))
-    out_I_float_1 = np_as_located_field(IDim)(np.zeros((size), dtype=float64))
+    out_I_float = np_as_located_field(IDim)(np.zeros((size,), dtype=float64))
+    out_I_float_1 = np_as_located_field(IDim)(np.zeros((size,), dtype=float64))
 
     @field_operator(backend=fieldview_backend)
     def multicopy(
@@ -76,7 +76,7 @@ def test_multicopy(fieldview_backend):
 
 def test_cartesian_shift(fieldview_backend):
     a = np_as_located_field(IDim)(np.arange(size + 1, dtype=np.float64))
-    out_I_float = np_as_located_field(IDim)(np.zeros((size), dtype=float64))
+    out_I_float = np_as_located_field(IDim)(np.zeros((size,), dtype=float64))
 
     @field_operator
     def shift_by_one(inp: Field[[IDim], float64]) -> Field[[IDim], float64]:
@@ -114,7 +114,7 @@ def test_fold_shifts(fieldview_backend):
     """Shifting the result of an addition should work."""
     a = np_as_located_field(IDim)(np.arange(size + 1, dtype=np.float64))
     b = np_as_located_field(IDim)(np.ones((size + 2)) * 2)
-    out_I_float = np_as_located_field(IDim)(np.zeros((size), dtype=float64))
+    out_I_float = np_as_located_field(IDim)(np.zeros((size,), dtype=float64))
 
     @field_operator
     def auto_lift(
@@ -137,7 +137,7 @@ def test_fold_shifts(fieldview_backend):
 def test_tuples(fieldview_backend):
     a_I_float = np_as_located_field(IDim)(np.random.randn(size).astype("float64"))
     b_I_float = np_as_located_field(IDim)(np.random.randn(size).astype("float64"))
-    out_I_float = np_as_located_field(IDim)(np.zeros((size), dtype=float64))
+    out_I_float = np_as_located_field(IDim)(np.zeros((size,), dtype=float64))
 
     @field_operator
     def tuples(
@@ -197,7 +197,7 @@ def test_scalar_arg_with_field(fieldview_backend):
 
     inp = index_field(Edge, dtype=float64)
     factor = 3.0
-    out = np_as_located_field(Edge)(np.zeros((size), dtype=np.float64))
+    out = np_as_located_field(Edge)(np.zeros((size,), dtype=np.float64))
 
     @field_operator
     def scalar_and_field_args(
@@ -287,7 +287,7 @@ def test_tuple_scalar_scan(fieldview_backend):
 
 def test_astype_int(fieldview_backend):
     size = 10
-    b_float_64 = np_as_located_field(IDim)(np.ones((size), dtype=np.float64))
+    b_float_64 = np_as_located_field(IDim)(np.full((size,), fill_value=1.5, dtype=np.float64))
     c_int64 = np_as_located_field(IDim)(np.ones((size,), dtype=np.int64))
     out_int_64 = np_as_located_field(IDim)(np.zeros((size,), dtype=np.int64))
 
@@ -301,7 +301,7 @@ def test_astype_int(fieldview_backend):
 
 
 def test_astype_bool(fieldview_backend):
-    b_float_64 = np_as_located_field(IDim)(np.ones((size), dtype=np.float64))
+    b_float_64 = np_as_located_field(IDim)(np.full((size,), fill_value=0.5, dtype=np.float64))
     c_bool = np_as_located_field(IDim)(np.ones((size,), dtype=bool))
     out_bool = np_as_located_field(IDim)(np.zeros((size,), dtype=bool))
 
@@ -315,17 +315,18 @@ def test_astype_bool(fieldview_backend):
 
 
 def test_astype_float(fieldview_backend):
-    c_int64 = np_as_located_field(IDim)(np.ones((size,), dtype=np.int64))
-    c_int32 = np_as_located_field(IDim)(np.ones((size,), dtype=np.int32))
-    out_int_32 = np_as_located_field(IDim)(np.zeros((size,), dtype=np.int32))
+    c_float64 = np_as_located_field(IDim)(
+        np.full((size,), fill_value=np.float64("5e300"), dtype=np.float64)
+    )
+    out_int_32 = np_as_located_field(IDim)(np.zeros((size,), dtype=np.float32))
 
     @field_operator(backend=fieldview_backend)
-    def astype_fieldop_float(b: Field[[IDim], int64]) -> Field[[IDim], int32]:
-        d = astype(b, int32)
+    def astype_fieldop_float(b: Field[[IDim], float64]) -> Field[[IDim], np.float32]:
+        d = astype(b, float32)
         return d
 
-    astype_fieldop_float(c_int64, out=out_int_32, offset_provider={})
-    assert np.allclose(c_int32.array(), out_int_32)
+    astype_fieldop_float(c_float64, out=out_int_32, offset_provider={})
+    assert np.all(out_int_32.array() == np.inf)
 
 
 def test_offset_field(fieldview_backend):
@@ -334,7 +335,7 @@ def test_offset_field(fieldview_backend):
     a_I_float_1 = np_as_located_field(IDim, KDim)(
         np.append(np.insert(a_I_arr, size, 0, axis=1), [np.array([0] * (size + 1))], axis=0)
     )
-    offset_field_arr = np.asarray(np.ones((size - 1, size - 1)), dtype=int64)
+    offset_field_arr = np.ones((size - 1, size - 1), dtype=int64)
     offset_field_comp = np.append(
         np.insert(offset_field_arr, size - 1, 0, axis=1), [np.array([0] * size)], axis=0
     )
@@ -375,7 +376,7 @@ def test_offset_field(fieldview_backend):
 def test_nested_tuple_return(fieldview_backend):
     a_I_float = np_as_located_field(IDim)(np.random.randn(size).astype("float64"))
     b_I_float = np_as_located_field(IDim)(np.random.randn(size).astype("float64"))
-    out_I_float = np_as_located_field(IDim)(np.zeros((size), dtype=float64))
+    out_I_float = np_as_located_field(IDim)(np.zeros((size,), dtype=float64))
 
     @field_operator
     def pack_tuple(
@@ -507,7 +508,7 @@ def test_tuple_with_local_field_in_reduction_shifted(reduction_setup, fieldview_
 def test_tuple_arg(fieldview_backend):
     a_I_float = np_as_located_field(IDim)(np.random.randn(size).astype("float64"))
     b_I_float = np_as_located_field(IDim)(np.random.randn(size).astype("float64"))
-    out_I_float = np_as_located_field(IDim)(np.zeros((size), dtype=float64))
+    out_I_float = np_as_located_field(IDim)(np.zeros((size,), dtype=float64))
 
     if fieldview_backend in [gtfn_cpu.run_gtfn, gtfn_cpu.run_gtfn_imperative]:
         pytest.skip("Tuple arguments are not supported in gtfn yet.")
@@ -586,16 +587,14 @@ def test_solve_triag(fieldview_backend):
 
     solve_tridiag(a, b, c, d, out=out, offset_provider={})
 
-    np.allclose(expected, out)
+    assert np.allclose(expected, out)
 
 
-def test_ternary_operator(fieldview_backend):
+@pytest.mark.parametrize("left,right", [(2.0, 3.0), (3.0, 2.0)])
+def test_ternary_operator(left, right, fieldview_backend):
     a_I_float = np_as_located_field(IDim)(np.random.randn(size).astype("float64"))
     b_I_float = np_as_located_field(IDim)(np.random.randn(size).astype("float64"))
-    out_I_float = np_as_located_field(IDim)(np.zeros((size), dtype=float64))
-
-    left = 2.0
-    right = 3.0
+    out_I_float = np_as_located_field(IDim)(np.zeros((size,), dtype=float64))
 
     @field_operator(backend=fieldview_backend)
     def ternary_field_op(
@@ -605,27 +604,25 @@ def test_ternary_operator(fieldview_backend):
 
     ternary_field_op(a_I_float, b_I_float, left, right, out=out_I_float, offset_provider={})
     e = np.asarray(a_I_float) if left < right else np.asarray(b_I_float)
-    np.allclose(e, out_I_float)
+    assert np.allclose(e, out_I_float)
 
     @field_operator(backend=fieldview_backend)
     def ternary_field_op_scalars(left: float, right: float) -> Field[[IDim], float]:
         return broadcast(3.0, (IDim,)) if left > right else broadcast(4.0, (IDim,))
 
     ternary_field_op_scalars(left, right, out=out_I_float, offset_provider={})
-    e = np.full(e.shape, 3.0) if left > right else e
-    np.allclose(e, out_I_float)
+    e = np.full(e.shape, 3.0) if left > right else np.full(e.shape, 4.0)
+    assert np.allclose(e, out_I_float)
 
 
-def test_ternary_operator_tuple(fieldview_backend):
+@pytest.mark.parametrize("left,right", [(2.0, 3.0), (3.0, 2.0)])
+def test_ternary_operator_tuple(left, right, fieldview_backend):
     if fieldview_backend in [gtfn_cpu.run_gtfn, gtfn_cpu.run_gtfn_imperative]:
         pytest.skip("Tuple arguments are not supported in gtfn yet.")
     a_I_float = np_as_located_field(IDim)(np.random.randn(size).astype("float64"))
     b_I_float = np_as_located_field(IDim)(np.random.randn(size).astype("float64"))
-    out_I_float = np_as_located_field(IDim)(np.zeros((size), dtype=float64))
-    out_I_float_1 = np_as_located_field(IDim)(np.zeros((size), dtype=float64))
-
-    left = 2.0
-    right = 3.0
+    out_I_float = np_as_located_field(IDim)(np.zeros((size,), dtype=float64))
+    out_I_float_1 = np_as_located_field(IDim)(np.zeros((size,), dtype=float64))
 
     @field_operator(backend=fieldview_backend)
     def ternary_field_op(
@@ -642,8 +639,8 @@ def test_ternary_operator_tuple(fieldview_backend):
         if left < right
         else (np.asarray(b_I_float), np.asarray(a_I_float))
     )
-    np.allclose(e, out_I_float)
-    np.allclose(f, out_I_float_1)
+    assert np.allclose(e, out_I_float)
+    assert np.allclose(f, out_I_float_1)
 
 
 def test_ternary_builtin_neighbor_sum(reduction_setup, fieldview_backend):
@@ -722,7 +719,7 @@ def test_scan_nested_tuple_input(fieldview_backend, forward):
         pytest.xfail("gtfn does not yet support scan pass or tuple arguments.")
 
     init = 1.0
-    inp1 = np_as_located_field(KDim)(np.ones(size))
+    inp1 = np_as_located_field(KDim)(np.ones((size,)))
     inp2 = np_as_located_field(KDim)(np.arange(0.0, size, 1))
     out = np_as_located_field(KDim)(np.zeros((size,)))
 
@@ -760,58 +757,61 @@ def test_docstring():
 
 
 def test_domain(fieldview_backend):
-    a_IJ_float = np_as_located_field(IDim, JDim)(np.ones((size, size), dtype=float64))
+    inp = np_as_located_field(IDim, JDim)(np.ones((size, size), dtype=float64))
+    out = np_as_located_field(IDim, JDim)(2 * np.ones((size, size), dtype=float64))
+
+    expected = np.array(out)
+    expected[1:9, 4:6] = 1 + 1
 
     @field_operator(backend=fieldview_backend)
     def fieldop_domain(a: Field[[IDim, JDim], float64]) -> Field[[IDim, JDim], float64]:
         return a + a
 
     @program
-    def program_domain(a: Field[[IDim, JDim], float64]):
-        fieldop_domain(a, out=a, domain={IDim: (minimum(1, 2), 9), JDim: (4, maximum(5, 6))})
+    def program_domain(inp: Field[[IDim, JDim], float64], out: Field[[IDim, JDim], float64]):
+        fieldop_domain(inp, out=out, domain={IDim: (minimum(1, 2), 9), JDim: (4, maximum(5, 6))})
 
-    program_domain(a_IJ_float, offset_provider={})
+    program_domain(inp, out, offset_provider={})
 
-    expected = np.asarray(a_IJ_float)
-    expected[1:9, 4:6] = 1 + 1
-
-    assert np.allclose(expected, a_IJ_float)
+    assert np.allclose(expected, out)
 
 
 def test_domain_input_bounds(fieldview_backend):
     if fieldview_backend in [gtfn_cpu.run_gtfn, gtfn_cpu.run_gtfn_imperative]:
         pytest.skip("FloorDiv not fully supported in gtfn.")
-    a_IJ_float = np_as_located_field(IDim, JDim)(np.ones((size, size), dtype=float64))
+    inp = np_as_located_field(IDim, JDim)(np.ones((size, size), dtype=float64))
+    out = np_as_located_field(IDim, JDim)(2 * np.ones((size, size), dtype=float64))
 
     lower_i = 1
     upper_i = 9
     lower_j = 4
     upper_j = 6
 
+    expected = np.array(out)
+    expected[lower_i:upper_i, lower_j:upper_j] = 1 + 1
+
     @field_operator(backend=fieldview_backend)
     def fieldop_domain(a: Field[[IDim, JDim], float64]) -> Field[[IDim, JDim], float64]:
         return a + a
 
     @program
     def program_domain(
-        a: Field[[IDim, JDim], float64],
+        inp: Field[[IDim, JDim], float64],
+        out: Field[[IDim, JDim], float64],
         lower_i: int64,
         upper_i: int64,
         lower_j: int64,
         upper_j: int64,
     ):
         fieldop_domain(
-            a,
-            out=a,
+            inp,
+            out=out,
             domain={IDim: (lower_i, upper_i // 1), JDim: (lower_j**1, upper_j)},
         )
 
-    program_domain(a_IJ_float, lower_i, upper_i, lower_j, upper_j, offset_provider={})
-
-    expected = np.asarray(a_IJ_float)
-    expected[1:9, 4:6] = 1 + 1
+    program_domain(inp, out, lower_i, upper_i, lower_j, upper_j, offset_provider={})
 
-    assert np.allclose(expected, a_IJ_float)
+    assert np.allclose(expected, out)
 
 
 def test_domain_input_bounds_1(fieldview_backend):
@@ -822,9 +822,12 @@ def test_domain_input_bounds_1(fieldview_backend):
     lower_j = 4
     upper_j = 6
 
+    expected = np.array(a_IJ_float)
+    expected[lower_i:upper_i, lower_j:upper_j] = 1 + 1
+
     @field_operator(backend=fieldview_backend)
     def fieldop_domain(a: Field[[IDim, JDim], float64]) -> Field[[IDim, JDim], float64]:
-        return a * a
+        return a + a
 
     @program
     def program_domain(
@@ -842,33 +845,39 @@ def test_domain_input_bounds_1(fieldview_backend):
 
     program_domain(a_IJ_float, lower_i, upper_i, lower_j, upper_j, offset_provider={})
 
-    expected = np.asarray(a_IJ_float)
-    expected[1:9, 4:6] = 2 * 2
-
     assert np.allclose(expected, a_IJ_float)
 
 
 def test_domain_tuple(fieldview_backend):
-    a_IJ_float = np_as_located_field(IDim, JDim)(np.ones((size, size), dtype=float64))
-    b2d_float = a_IJ_float
+    inp0 = np_as_located_field(IDim, JDim)(np.ones((size, size), dtype=float64))
+    inp1 = np_as_located_field(IDim, JDim)(2 * np.ones((size, size), dtype=float64))
+    out0 = np_as_located_field(IDim, JDim)(3 * np.ones((size, size), dtype=float64))
+    out1 = np_as_located_field(IDim, JDim)(4 * np.ones((size, size), dtype=float64))
+
+    expected0 = np.array(out0)
+    expected0[1:9, 4:6] = (np.asarray(inp0) + np.asarray(inp1))[1:9, 4:6]
+    expected1 = np.array(out1)
+    expected1[1:9, 4:6] = np.asarray(inp1)[1:9, 4:6]
 
     @field_operator(backend=fieldview_backend)
     def fieldop_domain_tuple(
-        a: Field[[IDim, JDim], float64]
+        a: Field[[IDim, JDim], float64], b: Field[[IDim, JDim], float64]
     ) -> tuple[Field[[IDim, JDim], float64], Field[[IDim, JDim], float64]]:
-        return (a + a, a)
+        return (a + b, b)
 
     @program
-    def program_domain_tuple(a: Field[[IDim, JDim], float64], b: Field[[IDim, JDim], float64]):
-        fieldop_domain_tuple(a, out=(b, a), domain={IDim: (1, 9), JDim: (4, 6)})
-
-    program_domain_tuple(a_IJ_float, b2d_float, offset_provider={})
+    def program_domain_tuple(
+        inp0: Field[[IDim, JDim], float64],
+        inp1: Field[[IDim, JDim], float64],
+        out0: Field[[IDim, JDim], float64],
+        out1: Field[[IDim, JDim], float64],
+    ):
+        fieldop_domain_tuple(inp0, inp1, out=(out0, out1), domain={IDim: (1, 9), JDim: (4, 6)})
 
-    expected = np.asarray(a_IJ_float)
-    expected[1:9, 4:6] = 1 + 1
+    program_domain_tuple(inp0, inp1, out0, out1, offset_provider={})
 
-    assert np.allclose(np.asarray(a_IJ_float), a_IJ_float)
-    assert np.allclose(expected, b2d_float)
+    assert np.allclose(expected0, out0)
+    assert np.allclose(expected1, out1)
 
 
 def test_where_k_offset(fieldview_backend):
@@ -943,11 +952,11 @@ def test_tuple_unpacking(fieldview_backend):
         pytest.skip("Tuple arguments are not supported in gtfn yet.")
 
     size = 10
-    inp = np_as_located_field(IDim)(np.ones((size)))
-    out1 = np_as_located_field(IDim)(np.ones((size)))
-    out2 = np_as_located_field(IDim)(np.ones((size)))
-    out3 = np_as_located_field(IDim)(np.ones((size)))
-    out4 = np_as_located_field(IDim)(np.ones((size)))
+    inp = np_as_located_field(IDim)(np.ones((size,)))
+    out1 = np_as_located_field(IDim)(np.ones((size,)))
+    out2 = np_as_located_field(IDim)(np.ones((size,)))
+    out3 = np_as_located_field(IDim)(np.ones((size,)))
+    out4 = np_as_located_field(IDim)(np.ones((size,)))
 
     @field_operator(backend=fieldview_backend)
     def unpack(
@@ -976,8 +985,8 @@ def test_tuple_unpacking_star_multi(fieldview_backend):
         pytest.skip("Tuple arguments are not supported in gtfn yet.")
 
     size = 10
-    inp = np_as_located_field(IDim)(np.ones((size)))
-    out = tuple(np_as_located_field(IDim)(np.ones(size) * i) for i in range(3 * 4))
+    inp = np_as_located_field(IDim)(np.ones((size,)))
+    out = tuple(np_as_located_field(IDim)(np.ones((size,)) * i) for i in range(3 * 4))
 
     OutType = tuple[
         Field[[IDim], float64],
diff --git a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_trace_shifts.py b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_trace_shifts.py
index 4ad4d01a..8a5da3df 100644
--- a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_trace_shifts.py
+++ b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_trace_shifts.py
@@ -12,6 +12,7 @@
 #
 # SPDX-License-Identifier: GPL-3.0-or-later
 
+from gt4py.next.ffront import itir_makers as im
 from gt4py.next.iterator import ir
 from gt4py.next.iterator.transforms.trace_shifts import ALL_NEIGHBORS, TraceShifts
 
@@ -120,15 +121,52 @@ def test_neighbors():
 
 def test_reduce():
     testee = ir.StencilClosure(
-        stencil=ir.FunCall(
-            fun=ir.SymRef(id="reduce"),
-            args=[ir.SymRef(id="plus"), ir.SymRef(id="init")],
+        # λ(inp) → reduce(plus, init)(·inp)
+        stencil=ir.Lambda(
+            params=[ir.Sym(id="inp")],
+            expr=ir.FunCall(
+                fun=ir.FunCall(
+                    fun=ir.SymRef(id="reduce"), args=[ir.SymRef(id="plus"), ir.SymRef(id="init")]
+                ),
+                args=[ir.FunCall(fun=ir.SymRef(id="deref"), args=[ir.SymRef(id="inp")])],
+            ),
         ),
         inputs=[ir.SymRef(id="inp")],
         output=ir.SymRef(id="out"),
         domain=ir.FunCall(fun=ir.SymRef(id="cartesian_domain"), args=[]),
     )
-    expected = dict()
+    expected = {"inp": [()]}
+
+    actual = dict()
+    TraceShifts().visit(testee, shifts=actual)
+    assert actual == expected
+
+
+def test_shifted_literal():
+    "Test shifting an applied lift of a stencil returning a constant / literal works."
+    testee = ir.StencilClosure(
+        # λ(x) → ·⟪Iₒ, 1ₒ⟫((↑(λ() → 1))())
+        stencil=im.lambda__("x")(im.deref_(im.shift_("I", 1)(im.lift_(im.lambda__()(1))()))),
+        inputs=[ir.SymRef(id="inp")],
+        output=ir.SymRef(id="out"),
+        domain=ir.FunCall(fun=ir.SymRef(id="cartesian_domain"), args=[]),
+    )
+    expected = {"inp": []}
+
+    actual = dict()
+    TraceShifts().visit(testee, shifts=actual)
+    assert actual == expected
+
+
+def test_tuple_get_on_closure_input():
+    testee = ir.StencilClosure(
+        # λ(x) → (·⟪Iₒ, 1ₒ⟫(x))[0]
+        stencil=im.lambda__("x")(im.tuple_get_(0, im.deref_(im.shift_("I", 1)("x")))),
+        inputs=[ir.SymRef(id="inp")],
+        output=ir.SymRef(id="out"),
+        domain=ir.FunCall(fun=ir.SymRef(id="cartesian_domain"), args=[]),
+    )
+    expected = {"inp": [(ir.OffsetLiteral(value="I"), ir.OffsetLiteral(value=1))]}
 
     actual = dict()
     TraceShifts().visit(testee, shifts=actual)