diff --git a/tests/next_tests/fixtures/past_common.py b/tests/next_tests/fixtures/past_common.py index 718db6c3a6..0b2681059c 100644 --- a/tests/next_tests/fixtures/past_common.py +++ b/tests/next_tests/fixtures/past_common.py @@ -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 diff --git a/tests/next_tests/integration_tests/cases.py b/tests/next_tests/integration_tests/cases.py index 78e6c62781..e99b4c7eec 100644 --- a/tests/next_tests/integration_tests/cases.py +++ b/tests/next_tests/integration_tests/cases.py @@ -48,12 +48,9 @@ E2VDim, Edge, IDim, - Ioff, JDim, - Joff, KDim, KHalfDim, - Koff, V2EDim, Vertex, exec_alloc_descriptor, 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 7640553e6a..2a2713a216 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 @@ -30,9 +30,6 @@ "IDim", "JDim", "KDim", - "Ioff", - "Joff", - "Koff", "Vertex", "Edge", "Cell", @@ -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") @@ -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", 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 647a94c8a2..ad8c93ea1b 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 @@ -37,10 +37,8 @@ E2VDim, Edge, IDim, - Ioff, JDim, KDim, - Koff, V2EDim, Vertex, cartesian_case, @@ -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)() @@ -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)})() @@ -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")() @@ -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, @@ -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 ) @@ -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)() diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_gt4py_builtins.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_gt4py_builtins.py index 1a1984a71b..0043881703 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_gt4py_builtins.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_gt4py_builtins.py @@ -21,11 +21,8 @@ V2E, Edge, IDim, - Ioff, JDim, - Joff, KDim, - Koff, V2EDim, Vertex, cartesian_case, @@ -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): @@ -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] @@ -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( diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_program.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_program.py index 1abaa47d03..1099ae4671 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_program.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_program.py @@ -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, @@ -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 diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_where.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_where.py index cd10c10437..2550692248 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_where.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_where.py @@ -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 @@ -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( diff --git a/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_program.py b/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_program.py index 09dc04acb1..2bac637e9b 100644 --- a/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_program.py +++ b/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_program.py @@ -26,7 +26,6 @@ I = gtx.Dimension("I") -Ioff = gtx.FieldOffset("Ioff", source=I, target=(I,)) @fundef @@ -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, diff --git a/tests/next_tests/integration_tests/multi_feature_tests/ffront_tests/test_icon_like_scan.py b/tests/next_tests/integration_tests/multi_feature_tests/ffront_tests/test_icon_like_scan.py index 6d7fd9df2b..df143eda54 100644 --- a/tests/next_tests/integration_tests/multi_feature_tests/ffront_tests/test_icon_like_scan.py +++ b/tests/next_tests/integration_tests/multi_feature_tests/ffront_tests/test_icon_like_scan.py @@ -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): @@ -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 @@ -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 @@ -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 @@ -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,) diff --git a/tests/next_tests/integration_tests/multi_feature_tests/ffront_tests/test_laplacian.py b/tests/next_tests/integration_tests/multi_feature_tests/ffront_tests/test_laplacian.py index 5c6bd5a54a..7a07299b18 100644 --- a/tests/next_tests/integration_tests/multi_feature_tests/ffront_tests/test_laplacian.py +++ b/tests/next_tests/integration_tests/multi_feature_tests/ffront_tests/test_laplacian.py @@ -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, ) diff --git a/tests/next_tests/unit_tests/embedded_tests/test_basic_program.py b/tests/next_tests/unit_tests/embedded_tests/test_basic_program.py index 9b4dae165a..a3e2c18e34 100644 --- a/tests/next_tests/unit_tests/embedded_tests/test_basic_program.py +++ b/tests/next_tests/unit_tests/embedded_tests/test_basic_program.py @@ -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 @@ -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) diff --git a/tests/next_tests/unit_tests/iterator_tests/test_type_inference.py b/tests/next_tests/unit_tests/iterator_tests/test_type_inference.py index 0dbb9357c3..45296cacb2 100644 --- a/tests/next_tests/unit_tests/iterator_tests/test_type_inference.py +++ b/tests/next_tests/unit_tests/iterator_tests/test_type_inference.py @@ -33,10 +33,8 @@ Edge, Cell, IDim, - Ioff, JDim, KDim, - Koff, V2EDim, Vertex, exec_alloc_descriptor,