Refactor tree_map and replace apply_to_primitive_constituents

src/gt4py/next/embedded/operators.py

@@ -13,15 +13,14 @@
 # SPDX-License-Identifier: GPL-3.0-or-later
 
 import dataclasses
-from types import ModuleType
 from typing import Any, Callable, Generic, Optional, ParamSpec, Sequence, TypeVar
 
-import numpy as np
-
 from gt4py import eve
 from gt4py._core import definitions as core_defs
-from gt4py.next import common, errors, utils
+from gt4py.next import common, errors, field_utils, utils
 from gt4py.next.embedded import common as embedded_common, context as embedded_context
+from gt4py.next.field_utils import get_array_ns
+from gt4py.next.type_system import type_translation
 
 
 _P = ParamSpec("_P")
@@ -64,8 +63,10 @@ def __call__(  # type: ignore[override]
             # even if the scan dimension is not in the input, we can scan over it
             out_domain = common.Domain(*out_domain, (scan_range))
 
-        xp = _get_array_ns(*all_args)
-        res = _construct_scan_array(out_domain, xp)(self.init)
+        xp = get_array_ns(*all_args)
+        res = field_utils.field_from_typespec(out_domain, xp)(
+            type_translation.from_value(self.init)
+        )
 
         def scan_loop(hpos: Sequence[common.NamedIndex]) -> None:
             acc: core_defs.ScalarT | tuple[core_defs.ScalarT | tuple, ...] = self.init
@@ -163,31 +164,6 @@ def _intersect_scan_args(
     )
 
 
-def _get_array_ns(
-    *args: core_defs.Scalar | common.Field | tuple[core_defs.Scalar | common.Field | tuple, ...],
-) -> ModuleType:
-    for arg in utils.flatten_nested_tuple(args):
-        if hasattr(arg, "array_ns"):
-            return arg.array_ns
-    return np
-
-
-def _construct_scan_array(
-    domain: common.Domain,
-    xp: ModuleType,  # TODO(havogt) introduce a NDArrayNamespace protocol
-) -> Callable[
-    [core_defs.Scalar | tuple[core_defs.Scalar | tuple, ...]],
-    common.MutableField | tuple[common.MutableField | tuple, ...],
-]:
-    @utils.tree_map
-    def impl(init: core_defs.Scalar) -> common.MutableField:
-        res = common._field(xp.empty(domain.shape, dtype=type(init)), domain=domain)
-        assert isinstance(res, common.MutableField)
-        return res
-
-    return impl
-
-
 def _tuple_assign_value(
     pos: Sequence[common.NamedIndex],
     target: common.MutableField | tuple[common.MutableField | tuple, ...],

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

@@ -833,12 +833,17 @@ def _visit_astype(self, node: foast.Call, **kwargs: Any) -> foast.Call:
                 f"Invalid call to 'astype': second argument must be a scalar type, got '{new_type}'.",
             )
 
-        return_type = type_info.apply_to_primitive_constituents(
-            value.type,
+        # return_type = type_info.apply_to_primitive_constituents(
+        #     value.type,
+        #     lambda primitive_type: with_altered_scalar_kind(
+        #         primitive_type, getattr(ts.ScalarKind, new_type.id.upper())
+        #     ),
+        # )
+        return_type = type_info.type_tree_map(
             lambda primitive_type: with_altered_scalar_kind(
                 primitive_type, getattr(ts.ScalarKind, new_type.id.upper())
-            ),
-        )
+            )
+        )(value.type)
         assert isinstance(return_type, (ts.TupleType, ts.ScalarType, ts.FieldType))
 
         return foast.Call(

src/gt4py/next/ffront/past_to_itir.py

@@ -340,7 +340,7 @@ def _construct_itir_domain_arg(
             domain_args.append(
                 itir.FunCall(
                     fun=itir.SymRef(id="named_range"),
-                    args=[itir.AxisLiteral(value=dim.value), lower, upper],
+                    args=[itir.AxisLiteral(value=dim.value, kind=dim.kind), lower, upper],
                 )
             )
             domain_args_kind.append(dim.kind)
@@ -349,11 +349,6 @@ def _construct_itir_domain_arg(
             domain_builtin = "cartesian_domain"
         elif self.grid_type == common.GridType.UNSTRUCTURED:
             domain_builtin = "unstructured_domain"
-            # for no good reason, the domain arguments for unstructured need to be in order (horizontal, vertical)
-            if domain_args_kind[0] == common.DimensionKind.VERTICAL:
-                assert len(domain_args) == 2
-                assert domain_args_kind[1] == common.DimensionKind.HORIZONTAL
-                domain_args[0], domain_args[1] = domain_args[1], domain_args[0]
         else:
             raise AssertionError()
 

src/gt4py/next/ffront/type_info.py

@@ -37,7 +37,9 @@ def promote_el(type_el: ts.TypeSpec) -> ts.TypeSpec:
             return ts.FieldType(dims=[], dtype=type_el)
         return type_el
 
-    return type_info.apply_to_primitive_constituents(type_, promote_el)
+    # return  type_info.apply_to_primitive_constituents(type_, promote_el)
+
+    return type_info.type_tree_map(promote_el)(type_)
 
 
 def promote_zero_dims(
@@ -308,6 +310,16 @@ def return_type_scanop(
         #  field
         [callable_type.axis],
     )
-    return type_info.apply_to_primitive_constituents(
-        carry_dtype, lambda arg: ts.FieldType(dims=promoted_dims, dtype=cast(ts.ScalarType, arg))
-    )
+
+    # return type_info.apply_to_primitive_constituents(
+    #     carry_dtype, lambda arg: ts.FieldType(dims=promoted_dims, dtype=cast(ts.ScalarType, arg))
+    # )
+    # @utils.tree_map(collection_type=ts.TupleType, result_collection_constructor=lambda x: ts.TupleType(types=[*x]))
+    # def tmp(x):
+    #     return ts.FieldType(dims=promoted_dims, dtype=x)
+
+    # return  tmp(carry_dtype)
+
+    return type_info.type_tree_map(
+        lambda arg: ts.FieldType(dims=promoted_dims, dtype=cast(ts.ScalarType, arg))
+    )(carry_dtype)

src/gt4py/next/field_utils.py

@@ -12,11 +12,40 @@
 #
 # SPDX-License-Identifier: GPL-3.0-or-later
 
+from types import ModuleType
+from typing import Callable
+
 import numpy as np
 
+from gt4py._core import definitions as core_defs
 from gt4py.next import common, utils
+from gt4py.next.type_system import type_specifications as ts, type_translation
 
 
 @utils.tree_map
 def asnumpy(field: common.Field | np.ndarray) -> np.ndarray:
     return field.asnumpy() if isinstance(field, common.Field) else field
+
+
+def field_from_typespec(
+    domain: common.Domain, xp: ModuleType
+) -> Callable[..., common.MutableField | tuple[common.MutableField | tuple, ...]]:
+    @utils.tree_map(collection_type=ts.TupleType, result_collection_constructor=tuple)
+    def impl(type_: ts.ScalarType) -> common.MutableField:
+        res = common._field(
+            xp.empty(domain.shape, dtype=xp.dtype(type_translation.as_dtype(type_).scalar_type)),
+            domain=domain,
+        )
+        assert isinstance(res, common.MutableField)
+        return res
+
+    return impl
+
+
+def get_array_ns(
+    *args: core_defs.Scalar | common.Field | tuple[core_defs.Scalar | common.Field | tuple, ...],
+) -> ModuleType:
+    for arg in utils.flatten_nested_tuple(args):
+        if hasattr(arg, "array_ns"):
+            return arg.array_ns
+    return np

src/gt4py/next/iterator/builtins.py

@@ -23,6 +23,11 @@ def __init__(self) -> None:
         super().__init__("Backend not selected")
 
 
+@builtin_dispatch
+def as_fieldop(*args):
+    raise BackendNotSelectedError()
+
+
 @builtin_dispatch
 def deref(*args):
     raise BackendNotSelectedError()
@@ -430,6 +435,7 @@ def bool(*args):  # noqa: A001 [builtin-variable-shadowing]
     "cartesian_domain",
     "unstructured_domain",
     "named_range",
+    "as_fieldop",
     *MATH_BUILTINS,
 }