feat[next]: Slicing field to 0d to return field not scalar

src/gt4py/next/embedded/nd_array_field.py

@@ -203,15 +203,13 @@ def remap(
 
     __call__ = remap  # type: ignore[assignment]
 
-    def restrict(self, index: common.AnyIndexSpec) -> common.Field | core_defs.ScalarT:
+    def restrict(self, index: common.AnyIndexSpec) -> common.Field:
         new_domain, buffer_slice = self._slice(index)
-
         new_buffer = self.ndarray[buffer_slice]
-        if len(new_domain) == 0:
-            # TODO: assert core_defs.is_scalar_type(new_buffer), new_buffer
-            return new_buffer  # type: ignore[return-value] # I don't think we can express that we return `ScalarT` here
-        else:
-            return self.__class__.from_array(new_buffer, domain=new_domain)
+        if new_domain.ndim == 0:
+            assert self.array_ns.dtype(new_buffer) in core_defs.SCALAR_TYPES
+            return self._scalar_to_field(new_buffer)  # type: ignore[return-value, arg-type]
+        return self.__class__.from_array(new_buffer, domain=new_domain)
 
     __getitem__ = restrict
 
@@ -304,6 +302,12 @@ def _slice(
         assert common.is_relative_index_sequence(slice_)
         return new_domain, slice_
 
+    def _scalar_to_field(self, value: core_defs.Scalar) -> np.ndarray:
+        if self.array_ns == cp:
+            return cp.asarray(value)
+        else:
+            return np.asarray(value)
+
 
 @dataclasses.dataclass(frozen=True)
 class NdArrayConnectivityField(  # type: ignore[misc] # for __ne__, __eq__
@@ -429,7 +433,7 @@ def inverse_image(
 
         return new_dims
 
-    def restrict(self, index: common.AnyIndexSpec) -> common.Field | core_defs.IntegralScalar:
+    def restrict(self, index: common.AnyIndexSpec) -> common.Field:
         cache_key = (id(self.ndarray), self.domain, index)
 
         if (restricted_connectivity := self._cache.get(cache_key, None)) is None:

src/gt4py/next/embedded/operators.py

@@ -187,8 +187,7 @@ def _tuple_at(
 ) -> core_defs.Scalar | tuple[core_defs.ScalarT | tuple, ...]:
     @utils.tree_map
     def impl(field: common.Field | core_defs.Scalar) -> core_defs.Scalar:
-        res = field[pos] if common.is_field(field) else field
-        res = res.item() if hasattr(res, "item") else res  # extract scalar value from array
+        res = field[pos].item() if common.is_field(field) else field  # type: ignore[union-attr]
         assert core_defs.is_scalar_type(res)
         return res
 

src/gt4py/next/ffront/func_to_foast.py

@@ -289,7 +289,7 @@ def visit_Subscript(self, node: ast.Subscript, **kwargs) -> foast.Subscript:
             index = self._match_index(node.slice)
         except ValueError:
             raise errors.DSLError(
-                self.get_location(node.slice), "eXpected an integral index."
+                self.get_location(node.slice), "Expected an integral index."
             ) from None
 
         return foast.Subscript(

src/gt4py/next/ffront/past_passes/type_deduction.py

@@ -217,12 +217,12 @@ def visit_Call(self, node: past.Call, **kwargs):
                         f"'{new_kwargs['out'].type}'."
                     )
             elif new_func.id in ["minimum", "maximum"]:
-                if new_args[0].type != new_args[1].type:
+                if arg_types[0] != arg_types[1]:
                     raise ValueError(
                         f"First and second argument in '{new_func.id}' must be of the same type."
-                        f"Got '{new_args[0].type}' and '{new_args[1].type}'."
+                        f"Got '{arg_types[0]}' and '{arg_types[1]}'."
                     )
-                return_type = new_args[0].type
+                return_type = arg_types[0]
             else:
                 raise AssertionError(
                     "Only calls to 'FieldOperator', 'ScanOperator' or 'minimum' and 'maximum' builtins allowed."

src/gt4py/next/ffront/past_to_itir.py

@@ -305,7 +305,7 @@ def _visit_stencil_call_out_arg(
     ) -> tuple[itir.Expr, itir.FunCall]:
         if isinstance(out_arg, past.Subscript):
             # as the ITIR does not support slicing a field we have to do a deeper
-            #  inspection of the PAST to emulate the behaviour
+            # inspection of the PAST to emulate the behaviour
             out_field_name: past.Name = out_arg.value
             return (
                 self._construct_itir_out_arg(out_field_name),
@@ -382,12 +382,11 @@ def visit_BinOp(self, node: past.BinOp, **kwargs) -> itir.FunCall:
         )
 
     def visit_Call(self, node: past.Call, **kwargs) -> itir.FunCall:
-        if node.func.id in ["maximum", "minimum"] and len(node.args) == 2:
+        if node.func.id in ["maximum", "minimum"]:
+            assert len(node.args) == 2
             return itir.FunCall(
                 fun=itir.SymRef(id=node.func.id),
                 args=[self.visit(node.args[0]), self.visit(node.args[1])],
             )
         else:
-            raise AssertionError(
-                "Only 'minimum' and 'maximum' builtins supported supported currently."
-            )
+            raise NotImplementedError("Only 'minimum', and 'maximum' builtins supported currently.")

src/gt4py/next/iterator/embedded.py

@@ -748,7 +748,10 @@ def _make_tuple(
         else:
             try:
                 data = field_or_tuple.field_getitem(named_indices)
-                return data
+                if core_defs.is_scalar_type(data):
+                    return data  # type: ignore[return-value] # type assessed in if
+                assert data.ndim == 0
+                return data.item()
             except embedded_exceptions.IndexOutOfBounds:
                 return _UNDEFINED
     else:
@@ -1086,12 +1089,12 @@ def remap(self, index_field: common.ConnectivityField | fbuiltins.FieldOffset) -
         # TODO can be implemented by constructing and ndarray (but do we know of which kind?)
         raise NotImplementedError()
 
-    def restrict(self, item: common.AnyIndexSpec) -> common.Field | core_defs.int32:
+    def restrict(self, item: common.AnyIndexSpec) -> common.Field:
         if common.is_absolute_index_sequence(item) and all(common.is_named_index(e) for e in item):  # type: ignore[arg-type] # we don't want to pollute the typing of `is_absolute_index_sequence` for this temporary code # fmt: off
             d, r = item[0]
             assert d == self._dimension
             assert isinstance(r, int)
-            return self.dtype.scalar_type(r)
+            return np.asarray(r)  # type: ignore[return-value] # Field is a superset
         # TODO set a domain...
         raise NotImplementedError()
 
@@ -1205,9 +1208,11 @@ def remap(self, index_field: common.ConnectivityField | fbuiltins.FieldOffset) -
         # TODO can be implemented by constructing and ndarray (but do we know of which kind?)
         raise NotImplementedError()
 
-    def restrict(self, item: common.AnyIndexSpec) -> common.Field | core_defs.ScalarT:
+    def restrict(self, item: common.AnyIndexSpec) -> common.Field:
         # TODO set a domain...
-        return self._value
+        if core_defs.is_scalar_type(self._value):
+            return np.asarray(self._value)  # type: ignore[return-value] # Field is a superset
+        return self._value  # type: ignore[return-value] # scalar type in if statement above
 
     __call__ = remap
     __getitem__ = restrict

tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py

@@ -320,6 +320,21 @@ def testee(inp: gtx.Field[[KDim], float]) -> gtx.Field[[KDim], float]:
     cases.verify(cartesian_case, testee, inp, out=out, ref=expected)
 
 
+def test_single_value_field(cartesian_case):
+    @gtx.field_operator
+    def testee_fo(a: cases.IKField) -> cases.IKField:
+        return a
+
+    @gtx.program
+    def testee_prog(a: cases.IKField):
+        testee_fo(a, out=a[1:2, 3:4])
+
+    a = cases.allocate(cartesian_case, testee_prog, "a")()
+    ref = a[1, 3]
+
+    cases.verify(cartesian_case, testee_prog, a, inout=a[1, 3], ref=ref)
+
+
 def test_astype_int(cartesian_case):  # noqa: F811 # fixtures
     @gtx.field_operator
     def testee(a: cases.IFloatField) -> gtx.Field[[IDim], int64]:

tests/next_tests/integration_tests/feature_tests/iterator_tests/test_horizontal_indirection.py

@@ -70,7 +70,7 @@ def test_simple_indirection(program_processor):
 
     ref = np.zeros(shape, dtype=inp.dtype)
     for i in range(shape[0]):
-        ref[i] = inp.ndarray[i + 1 - 1] if cond[i] < 0.0 else inp.ndarray[i + 1 + 1]
+        ref[i] = inp.asnumpy()[i + 1 - 1] if cond.asnumpy()[i] < 0.0 else inp.asnumpy()[i + 1 + 1]
 
     run_processor(
         conditional_indirection[cartesian_domain(named_range(IDim, 0, shape[0]))],
@@ -101,7 +101,7 @@ def test_direct_offset_for_indirection(program_processor):
 
     ref = np.zeros(shape)
     for i in range(shape[0]):
-        ref[i] = inp[i + cond[i]]
+        ref[i] = inp.asnumpy()[i + cond.asnumpy()[i]]
 
     run_processor(
         direct_indirection[cartesian_domain(named_range(IDim, 0, shape[0]))],

tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_anton_toy.py

@@ -65,11 +65,15 @@ def fencil(x, y, z, out, inp):
 def naive_lap(inp):
     shape = [inp.shape[0] - 2, inp.shape[1] - 2, inp.shape[2]]
     out = np.zeros(shape)
+    inp_data = inp.asnumpy()
     for i in range(1, shape[0] + 1):
         for j in range(1, shape[1] + 1):
             for k in range(0, shape[2]):
-                out[i - 1, j - 1, k] = -4 * inp[i, j, k] + (
-                    inp[i + 1, j, k] + inp[i - 1, j, k] + inp[i, j + 1, k] + inp[i, j - 1, k]
+                out[i - 1, j - 1, k] = -4 * inp_data[i, j, k] + (
+                    inp_data[i + 1, j, k]
+                    + inp_data[i - 1, j, k]
+                    + inp_data[i, j + 1, k]
+                    + inp_data[i, j - 1, k]
                 )
     return out
 

tests/next_tests/unit_tests/embedded_tests/test_nd_array_field.py

@@ -468,10 +468,11 @@ def test_absolute_indexing_value_return():
     field = common._field(np.reshape(np.arange(100, dtype=np.int32), (10, 10)), domain=domain)
 
     named_index = ((IDim, 12), (JDim, 6))
+    assert common.is_field(field)
     value = field[named_index]
 
-    assert isinstance(value, np.int32)
-    assert value == 21
+    assert isinstance(value, np.ndarray)
+    assert value == np.asarray(21)
 
 
 @pytest.mark.parametrize(
@@ -568,14 +569,17 @@ def test_relative_indexing_slice_3D(index, expected_shape, expected_domain):
 
 @pytest.mark.parametrize(
     "index, expected_value",
-    [((1, 0), 10), ((0, 1), 1)],
+    [
+        ((1, 0), 10),
+        ((0, 1), 1),
+    ],
 )
 def test_relative_indexing_value_return(index, expected_value):
     domain = common.Domain(dims=(IDim, JDim), ranges=(UnitRange(5, 15), UnitRange(2, 12)))
     field = common._field(np.reshape(np.arange(100, dtype=int), (10, 10)), domain=domain)
     indexed_field = field[index]
 
-    assert indexed_field == expected_value
+    assert indexed_field.item() == expected_value
 
 
 @pytest.mark.parametrize("lazy_slice", [lambda f: f[13], lambda f: f[:5, :3, :2]])