removed [0][0] indexing

GridTools · Jan 5, 2024 · 90f6796 · 90f6796
1 parent 9f2bcc6
commit 90f6796
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/src/gt4py/next/embedded/nd_array_field.py b/src/gt4py/next/embedded/nd_array_field.py
@@ -606,7 +606,9 @@ def _astype(field: common.Field | core_defs.ScalarT | tuple, type_: type) -> NdA
 def _as_offset(offset_: fbuiltins.FieldOffset, field: common.Field) -> NdArrayConnectivityField:
     if isinstance(field, NdArrayField):
         # change field.ndarray from relative to absolute
-        offset_dim = np.where(list(map(lambda x: x == offset_.source, field.domain.dims)))[0][0]
+        offset_dim = np.squeeze(
+            np.where(list(map(lambda x: x == offset_.source, field.domain.dims)))
+        ).item()
         new_connectivity = np.indices(field.ndarray.shape)[offset_dim] + field.ndarray
         return NumPyArrayConnectivityField.from_array(
             new_connectivity, codomain=offset_.source, domain=field.domain

diff --git a/src/gt4py/next/ffront/experimental.py b/src/gt4py/next/ffront/experimental.py
@@ -24,7 +24,9 @@ def as_offset(
     field: common.Field,
     /,
 ) -> common.ConnectivityField:
-    offset_dim = np.where(list(map(lambda x: x == offset_.source, field.domain.dims)))[0][0]
+    offset_dim = np.squeeze(
+        np.where(list(map(lambda x: x == offset_.source, field.domain.dims)))
+    ).item()
     new_connectivity = np.indices(field.ndarray.shape)[offset_dim] + field.ndarray
     return common.connectivity(new_connectivity, codomain=offset_.source, domain=field.domain)