Skip to content

refactor[next]: Use cartesian shift syntax in all tests #2411

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
tehrengruber wants to merge 2 commits into
base: main
Choose a base branch
from
Draft

refactor[next]: Use cartesian shift syntax in all tests #2411

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
491d1c2
Use cartesian shift syntax in all tests
tehrengruber Dec 4, 2025
b69e966
Remove default offset provider and offsets from default tests
tehrengruber Dec 4, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 0 additions & 2 deletions tests/next_tests/fixtures/past_common.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,9 +15,7 @@


IDim = gtx.Dimension("IDim")
Ioff = gtx.FieldOffset("Ioff", source=IDim, target=(IDim,))
JDim = gtx.Dimension("JDim")
Joff = gtx.FieldOffset("Joff", source=JDim, target=(JDim,))


# TODO(tehrengruber): Improve test structure. Identity needs to be decorated
Expand Down
3 changes: 0 additions & 3 deletions tests/next_tests/integration_tests/cases.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,12 +48,9 @@
E2VDim,
Edge,
IDim,
Ioff,
JDim,
Joff,
KDim,
KHalfDim,
Koff,
V2EDim,
Vertex,
exec_alloc_descriptor,
Expand Down
13 changes: 1 addition & 12 deletions tests/next_tests/integration_tests/feature_tests/ffront_tests/ffront_test_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,9 +30,6 @@
"IDim",
"JDim",
"KDim",
"Ioff",
"Joff",
"Koff",
"Vertex",
"Edge",
"Cell",
Expand Down Expand Up @@ -136,9 +133,6 @@ def debug_itir(tree):
JDim = gtx.Dimension("JDim")
KDim = gtx.Dimension("KDim", kind=gtx.DimensionKind.VERTICAL)
KHalfDim = gtx.Dimension("KHalf", kind=gtx.DimensionKind.VERTICAL)
Ioff = gtx.FieldOffset("Ioff", source=IDim, target=(IDim,))
Joff = gtx.FieldOffset("Joff", source=JDim, target=(JDim,))
Koff = gtx.FieldOffset("Koff", source=KDim, target=(KDim,))

Vertex = gtx.Dimension("Vertex")
Edge = gtx.Dimension("Edge")
Expand Down Expand Up @@ -178,12 +172,7 @@ def simple_cartesian_grid(
sizes = (sizes,) * 4
assert len(sizes) == 4, "sizes must be a tuple of four integers"

offset_provider = {
"Ioff": IDim,
"Joff": JDim,
"Koff": KDim,
"KHalfoff": KHalfDim,
}
offset_provider = {}

return types.SimpleNamespace(
name="simple_cartesian_grid",
Expand Down
19 changes: 10 additions & 9 deletions tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,10 +37,8 @@
E2VDim,
Edge,
IDim,
Ioff,
JDim,
KDim,
Koff,
V2EDim,
Vertex,
cartesian_case,
Expand Down Expand Up @@ -119,7 +117,7 @@ def testee() -> cases.IFloatField:
def test_cartesian_shift(cartesian_case):
@gtx.field_operator
def testee(a: cases.IJKField) -> cases.IJKField:
return a(Ioff[1])
return a(IDim + 1)

a = cases.allocate(cartesian_case, testee, "a").extend({IDim: (0, 1)})()
out = cases.allocate(cartesian_case, testee, cases.RETURN)()
Expand Down Expand Up @@ -205,8 +203,8 @@ def test_fold_shifts(cartesian_case):

@gtx.field_operator
def testee(a: cases.IJKField, b: cases.IJKField) -> cases.IJKField:
tmp = a + b(Ioff[1])
return tmp(Ioff[1])
tmp = a + b(IDim + 1)
return tmp(IDim + 1)

a = cases.allocate(cartesian_case, testee, "a").extend({cases.IDim: (0, 1)})()
b = cases.allocate(cartesian_case, testee, "b").extend({cases.IDim: (0, 2)})()
Expand Down Expand Up @@ -355,7 +353,7 @@ def test_scalar_arg_with_field(cartesian_case):
@gtx.field_operator
def testee(a: cases.IJKField, b: int32) -> cases.IJKField:
tmp = b * a
return tmp(Ioff[1])
return tmp(IDim + 1)

a = cases.allocate(cartesian_case, testee, "a").extend({IDim: (0, 1)})()
b = cases.allocate(cartesian_case, testee, "b")()
Expand Down Expand Up @@ -452,7 +450,7 @@ def testee_scan(state: float, inp: float) -> float:

@gtx.field_operator
def testee(inp: gtx.Field[[KDim], float]) -> gtx.Field[[KDim], float]:
return testee_scan(inp(Koff[1]))
return testee_scan(inp(KDim + 1))

inp = cases.allocate(
cartesian_case,
Expand Down Expand Up @@ -646,6 +644,9 @@ def testee() -> gtx.Field[[IDim], int_alias]:

@pytest.mark.uses_dynamic_offsets
def test_offset_field(cartesian_case):
Ioff = gtx.FieldOffset("Ioff", source=IDim, target=(IDim,))
Koff = gtx.FieldOffset("Koff", source=KDim, target=(KDim,))

ref = np.full(
(cartesian_case.default_sizes[IDim], cartesian_case.default_sizes[KDim]), True, dtype=bool
)
Expand All @@ -656,8 +657,8 @@ def testee(a: cases.IKField, offset_field: cases.IKField) -> gtx.Field[[IDim, KD
# note: this leads to an access to offset_field in
# IDim: (0, out.size[I]), KDim: (0, out.size[K]+1)
a_i_k = a_i(as_offset(Koff, offset_field))
b_i = a(Ioff[1])
b_i_k = b_i(Koff[1])
b_i = a(IDim + 1)
b_i_k = b_i(KDim + 1)
return a_i_k == b_i_k

out = cases.allocate(cartesian_case, testee, cases.RETURN)()
Expand Down
9 changes: 3 additions & 6 deletions tests/next_tests/integration_tests/feature_tests/ffront_tests/test_gt4py_builtins.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,11 +21,8 @@
V2E,
Edge,
IDim,
Ioff,
JDim,
Joff,
KDim,
Koff,
V2EDim,
Vertex,
cartesian_case,
Expand Down Expand Up @@ -146,7 +143,7 @@ def reduction(edge_f: EKField) -> VKField:
@gtx.field_operator
def fencil_op(edge_f: EKField) -> VKField:
red = reduction(edge_f)
return red(Koff[1])
return red(KDim + 1)

@gtx.program
def fencil(edge_f: EKField, out: VKField):
Expand Down Expand Up @@ -377,7 +374,7 @@ def test_broadcast_shifted(cartesian_case):
@gtx.field_operator
def simple_broadcast(inp: cases.IField) -> cases.IJField:
bcasted = broadcast(inp, (IDim, JDim))
return bcasted(Joff[1])
return bcasted(JDim + 1)

cases.verify_with_default_data(
cartesian_case, simple_broadcast, ref=lambda inp: inp[:, np.newaxis]
Expand Down Expand Up @@ -439,7 +436,7 @@ def conditional_shifted(
mask: cases.IBoolField, a: cases.IFloatField, b: cases.IFloatField
) -> gtx.Field[[IDim], float64]:
tmp = where(mask, a, b)
return tmp(Ioff[1])
return tmp(IDim + 1)

@gtx.program
def conditional_program(
Expand Down
3 changes: 1 addition & 2 deletions tests/next_tests/integration_tests/feature_tests/ffront_tests/test_program.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@
from next_tests.integration_tests import cases
from next_tests.integration_tests.cases import (
IDim,
Ioff,
JDim,
cartesian_case,
exec_alloc_descriptor,
Expand Down Expand Up @@ -56,7 +55,7 @@ def test_identity_fo_execution(cartesian_case, identity_def):
def test_shift_by_one_execution(cartesian_case):
@gtx.field_operator
def shift_by_one(in_field: cases.IFloatField) -> cases.IFloatField:
return in_field(Ioff[1])
return in_field(IDim + 1)

# direct call to field operator
# TODO(tehrengruber): slicing located fields not supported currently
Expand Down
4 changes: 2 additions & 2 deletions tests/next_tests/integration_tests/feature_tests/ffront_tests/test_where.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
import numpy as np
from typing import Tuple
import pytest
from next_tests.integration_tests.cases import IDim, JDim, KDim, Koff, cartesian_case
from next_tests.integration_tests.cases import IDim, JDim, KDim, cartesian_case
from gt4py import next as gtx
from gt4py.next import int32
from gt4py.next.ffront.fbuiltins import where, broadcast
Expand All @@ -25,7 +25,7 @@ def test_where_k_offset(cartesian_case):
def fieldop_where_k_offset(
inp: cases.IKField, k_index: gtx.Field[[KDim], gtx.IndexType]
) -> cases.IKField:
return where(k_index > 0, inp(Koff[-1]), 2)
return where(k_index > 0, inp(KDim - 1), 2)

@gtx.program
def prog(
Expand Down
4 changes: 2 additions & 2 deletions tests/next_tests/integration_tests/feature_tests/iterator_tests/test_program.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,6 @@


I = gtx.Dimension("I")
Ioff = gtx.FieldOffset("Ioff", source=I, target=(I,))


@fundef
Expand Down Expand Up @@ -80,7 +79,8 @@ def test_index_builtin(program_processor):
def index_program_shift(out, size):
set_at(
as_fieldop(
lambda i: deref(i) + deref(shift(Ioff, 1)(i)), cartesian_domain(named_range(I, 0, size))
lambda i: deref(i) + deref(shift("Ioff", 1)(i)),
cartesian_domain(named_range(I, 0, size)),
)(index(I)),
cartesian_domain(named_range(I, 0, size)),
out,
Expand Down
25 changes: 12 additions & 13 deletions tests/next_tests/integration_tests/multi_feature_tests/ffront_tests/test_icon_like_scan.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,6 @@

Cell = gtx.Dimension("Cell")
KDim = gtx.Dimension("KDim", kind=gtx.DimensionKind.VERTICAL)
Koff = gtx.FieldOffset("Koff", KDim, (KDim,))


class State(NamedTuple):
Expand Down Expand Up @@ -59,9 +58,9 @@ def _solve_nonhydro_stencil_52_like(
gtx.Field[[Cell, KDim], float], gtx.Field[[Cell, KDim], float], gtx.Field[[Cell, KDim], bool]
]:
"""No projector required as we write all output of the scan (including dummy field)"""
z_a = z_beta(Koff[-1]) * z_alpha(Koff[-1])
z_c = z_beta * z_alpha(Koff[1])
z_b = z_alpha * (z_beta(Koff[-1]) + z_beta)
z_a = z_beta(KDim - 1) * z_alpha(KDim - 1)
z_c = z_beta * z_alpha(KDim + 1)
z_b = z_alpha * (z_beta(KDim - 1) + z_beta)
z_q_res, w_res, dummy = _scan(w, z_q, z_a, z_b, z_c)
return z_q_res, w_res, dummy

Expand All @@ -86,9 +85,9 @@ def _solve_nonhydro_stencil_52_like_with_gtfn_tuple_merge(
z_q: gtx.Field[[Cell, KDim], float],
w: gtx.Field[[Cell, KDim], float],
) -> tuple[gtx.Field[[Cell, KDim], float], gtx.Field[[Cell, KDim], float]]:
z_a = z_beta(Koff[-1]) * z_alpha(Koff[-1])
z_c = z_beta * z_alpha(Koff[1])
z_b = z_alpha * (z_beta(Koff[-1]) + z_beta)
z_a = z_beta(KDim - 1) * z_alpha(KDim - 1)
z_c = z_beta * z_alpha(KDim + 1)
z_b = z_alpha * (z_beta(KDim - 1) + z_beta)
z_q_res, w_res, _ = _scan(w, z_q, z_a, z_b, z_c)
return z_q_res, w_res

Expand All @@ -112,9 +111,9 @@ def _solve_nonhydro_stencil_52_like_z_q(
z_q: gtx.Field[[Cell, KDim], float],
w: gtx.Field[[Cell, KDim], float],
) -> gtx.Field[[Cell, KDim], float]:
z_a = z_beta(Koff[-1]) * z_alpha(Koff[-1])
z_c = z_beta * z_alpha(Koff[1])
z_b = z_alpha * (z_beta(Koff[-1]) + z_beta)
z_a = z_beta(KDim - 1) * z_alpha(KDim - 1)
z_c = z_beta * z_alpha(KDim + 1)
z_b = z_alpha * (z_beta(KDim - 1) + z_beta)
z_q_res, w_res, _ = _scan(w, z_q, z_a, z_b, z_c)
return z_q_res

Expand All @@ -137,9 +136,9 @@ def _solve_nonhydro_stencil_52_like_z_q_tup(
z_q: gtx.Field[[Cell, KDim], float],
w: gtx.Field[[Cell, KDim], float],
) -> tuple[gtx.Field[[Cell, KDim], float]]:
z_a = z_beta(Koff[-1]) * z_alpha(Koff[-1])
z_c = z_beta * z_alpha(Koff[1])
z_b = z_alpha * (z_beta(Koff[-1]) + z_beta)
z_a = z_beta(KDim - 1) * z_alpha(KDim - 1)
z_c = z_beta * z_alpha(KDim + 1)
z_b = z_alpha * (z_beta(KDim - 1) + z_beta)
z_q_res, w_res, _ = _scan(w, z_q, z_a, z_b, z_c)
return (z_q_res,)

Expand Down
2 changes: 1 addition & 1 deletion tests/next_tests/integration_tests/multi_feature_tests/ffront_tests/test_laplacian.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
import gt4py.next as gtx

from next_tests.integration_tests import cases
from next_tests.integration_tests.cases import IDim, JDim, Joff, cartesian_case
from next_tests.integration_tests.cases import IDim, JDim, cartesian_case
from next_tests.integration_tests.feature_tests.ffront_tests.ffront_test_utils import (
exec_alloc_descriptor,
)
Expand Down
5 changes: 2 additions & 3 deletions tests/next_tests/unit_tests/embedded_tests/test_basic_program.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,14 +12,13 @@


IDim = gtx.Dimension("IDim")
IOff = gtx.FieldOffset("IOff", source=IDim, target=(IDim,))


@gtx.field_operator
def fop(
a: gtx.Field[[IDim], gtx.float64], b: gtx.Field[[IDim], gtx.float64]
) -> gtx.Field[[IDim], gtx.float64]:
return a(IOff[1]) + b
return a(IDim + 1) + b


@gtx.program
Expand All @@ -36,6 +35,6 @@ def test_basic():
b = gtx.as_field([(IDim, gtx.common.UnitRange(0, 4))], np.asarray([0.0, 1.0, 2.0, 3.0]))
out = gtx.as_field([(IDim, gtx.common.UnitRange(0, 4))], np.asarray([0.0, 0.0, 0.0, 0.0]))

prog(a, b, out, offset_provider={"IOff": IDim})
prog(a, b, out)
assert out.domain == b.domain
assert np.allclose(out.ndarray, a.ndarray + b.ndarray)
2 changes: 0 additions & 2 deletions tests/next_tests/unit_tests/iterator_tests/test_type_inference.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,10 +33,8 @@
Edge,
Cell,
IDim,
Ioff,
JDim,
KDim,
Koff,
V2EDim,
Vertex,
exec_alloc_descriptor,
Expand Down
Loading