From 7fbf5ad58211dc1d0f2343dddfd459c14976a15c Mon Sep 17 00:00:00 2001
From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com>
Date: Mon, 25 Nov 2024 09:02:44 +0100
Subject: [PATCH 1/2] small kdim fix

---
 .../next/program_processors/codegens/gtfn/itir_to_gtfn_ir.py   | 3 +++
 .../feature_tests/ffront_tests/ffront_test_utils.py            | 2 ++
 2 files changed, 5 insertions(+)

diff --git a/src/gt4py/next/program_processors/codegens/gtfn/itir_to_gtfn_ir.py b/src/gt4py/next/program_processors/codegens/gtfn/itir_to_gtfn_ir.py
index bc2bd645e8..6ad5066a03 100644
--- a/src/gt4py/next/program_processors/codegens/gtfn/itir_to_gtfn_ir.py
+++ b/src/gt4py/next/program_processors/codegens/gtfn/itir_to_gtfn_ir.py
@@ -198,6 +198,9 @@ def _collect_offset_definitions(
                         "Mapping an offset to a horizontal dimension in unstructured is not allowed."
                     )
                 # create alias from vertical offset to vertical dimension
+                offset_definitions[dim.value] = TagDefinition(
+                    name=Sym(id=dim.value), alias=_vertical_dimension
+                )
                 offset_definitions[offset_name] = TagDefinition(
                     name=Sym(id=offset_name), alias=SymRef(id=dim.value)
                 )
diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/ffront_test_utils.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/ffront_test_utils.py
index c64efb27d2..545b18bc21 100644
--- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/ffront_test_utils.py
+++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/ffront_test_utils.py
@@ -229,6 +229,7 @@ def simple_mesh() -> MeshDescriptor:
             C2E.value: gtx.NeighborTableOffsetProvider(
                 c2e_arr, Cell, Edge, 4, has_skip_values=False
             ),
+            Koff.value: KDim,
         },
     )
 
@@ -305,6 +306,7 @@ def skip_value_mesh() -> MeshDescriptor:
             C2E.value: gtx.NeighborTableOffsetProvider(
                 c2e_arr, Cell, Edge, 3, has_skip_values=False
             ),
+            Koff.value: KDim,
         },
     )
 

From 1714fb80e31e2d236e2bc82c016a9d5b15278b09 Mon Sep 17 00:00:00 2001
From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com>
Date: Mon, 25 Nov 2024 13:21:07 +0100
Subject: [PATCH 2/2] added execution test

---
 .../feature_tests/ffront_tests/test_execution.py  | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

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 a5453151e6..d1e8f4fb04 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
@@ -81,7 +81,7 @@ def testee(a: cases.IJKField) -> cases.IJKField:
 
 
 @pytest.mark.uses_unstructured_shift
-def test_unstructured_shift(unstructured_case):
+def test_unstructured_shift_indexed(unstructured_case):
     @gtx.field_operator
     def testee(a: cases.VField) -> cases.EField:
         return a(E2V[0])
@@ -93,6 +93,19 @@ def testee(a: cases.VField) -> cases.EField:
     )
 
 
+@pytest.mark.uses_unstructured_shift
+def test_unstructured_shift(unstructured_case):
+    @gtx.field_operator
+    def testee(a: cases.VField) -> cases.EField:
+        return neighbor_sum(a(E2V), axis=E2VDim)
+
+    cases.verify_with_default_data(
+        unstructured_case,
+        testee,
+        ref=lambda a: np.sum(a[unstructured_case.offset_provider["E2V"].table], axis=1),
+    )
+
+
 @pytest.mark.uses_unstructured_shift
 def test_composed_unstructured_shift(unstructured_case):
     @gtx.field_operator