edits

GridTools · Jan 15, 2024 · 2ea83f1 · 2ea83f1
1 parent 7c3e9cb
commit 2ea83f1
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 21 deletions.
diff --git a/src/gt4py/next/embedded/nd_array_field.py b/src/gt4py/next/embedded/nd_array_field.py
@@ -214,12 +214,10 @@ def remap(
 
             # then compute the index array
             xp = self.array_ns
+            new_idx_array = xp.asarray(restricted_connectivity.ndarray) - current_range.start
             if restricted_connectivity.domain.ndim > 1:
-                new_buffer = _take_mdim(
-                    self._ndarray, restricted_connectivity.ndarray, new_domain, dim
-                )
+                new_buffer = _take_mdim(self._ndarray, new_idx_array, new_domain, dim)
             else:
-                new_idx_array = xp.asarray(restricted_connectivity.ndarray) - current_range.start
                 # finally, take the new array
                 new_buffer = xp.take(self._ndarray, new_idx_array, axis=dim_idx)
 

diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_as_offset.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_as_offset.py
@@ -21,15 +21,7 @@
 import gt4py.next as gtx
 from gt4py.next.ffront.experimental import as_offset
 from tests.next_tests.integration_tests import cases
-from tests.next_tests.integration_tests.cases import (
-    IDim,
-    Ioff,
-    JDim,
-    KDim,
-    Koff,
-    cartesian_case,
-    unstructured_case,
-)
+from tests.next_tests.integration_tests.cases import IDim, Ioff, JDim, KDim, Koff, cartesian_case
 from tests.next_tests.integration_tests.feature_tests.ffront_tests.ffront_test_utils import (
     fieldview_backend,
     reduction_setup,
@@ -103,6 +95,9 @@ def testee(
 
 @pytest.mark.uses_dynamic_offsets
 def test_offset_field_input_domain(cartesian_case):
+    i_size = cartesian_case.default_sizes[IDim]
+    k_size = cartesian_case.default_sizes[KDim]
+
     @gtx.field_operator
     def testee_fo(
         a: gtx.Field[[IDim, KDim], int], offset_field: gtx.Field[[IDim, KDim], int]
@@ -117,14 +112,10 @@ def testee(
     ):
         testee_fo(a, offset_field, out=out)
 
-    # out = cases.allocate(cartesian_case, testee, "out")()
     a_fo = cases.allocate(cartesian_case, testee, "a").extend({IDim: (0, 3)})().ndarray
-    # a = gtx.as_field({IDim: (1, 5), KDim: (0, 10)}, a_fo)
     out = gtx.as_field(
         [IDim, KDim],
-        np.zeros(
-            [cartesian_case.default_sizes[IDim], cartesian_case.default_sizes[KDim]], dtype=int
-        ),
+        np.zeros([i_size, k_size], dtype=int),
         origin={IDim: 2, KDim: 0},
     )
     a = gtx.as_field([IDim, KDim], a_fo, origin={IDim: 2, KDim: 0})
@@ -137,11 +128,10 @@ def testee(
         testee,
         a,
         offset_field,
-        out,
-        inout=out,
+        out[2:],
+        inout=out[2:],
         offset_provider={"Ioff": IDim},
         ref=ref,
-        comparison=lambda out, ref: np.all(out == ref),
     )