Refactor tree_map and replace apply_to_primitive_constituents

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

@@ -848,12 +848,12 @@ 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(
+        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,
-        )
+            )
+        )(value.type)
+
         assert isinstance(return_type, (ts.TupleType, ts.ScalarType, ts.FieldType))
 
         return foast.Call(

src/gt4py/next/ffront/type_info.py

@@ -31,7 +31,7 @@ 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(promote_el, type_)
+    return type_info.type_tree_map(promote_el)(type_)
 
 
 def promote_zero_dims(
@@ -306,6 +306,7 @@ def return_type_scanop(
         #  field
         [callable_type.axis],
     )
-    return type_info.apply_to_primitive_constituents(
-        lambda arg: ts.FieldType(dims=promoted_dims, dtype=cast(ts.ScalarType, arg)), 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/iterator/transforms/fuse_as_fieldop.py

@@ -139,7 +139,7 @@ def fuse_as_fieldop(
             # just a safety check if typing information is available
             if arg.type and not isinstance(arg.type, ts.DeferredType):
                 assert isinstance(arg.type, ts.TypeSpec)
-                dtype = type_info.apply_to_primitive_constituents(type_info.extract_dtype, arg.type)
+                dtype = type_info.type_tree_map(type_info.extract_dtype)(arg.type)
                 assert not isinstance(dtype, it_ts.ListType)
             new_param: str
             if isinstance(
@@ -233,7 +233,7 @@ def visit_FunCall(self, node: itir.FunCall):
             eligible_args = []
             for arg, arg_shifts in zip(args, shifts, strict=True):
                 assert isinstance(arg.type, ts.TypeSpec)
-                dtype = type_info.apply_to_primitive_constituents(type_info.extract_dtype, arg.type)
+                dtype = type_info.type_tree_map(type_info.extract_dtype)(arg.type)
                 # TODO(tehrengruber): make this configurable
                 eligible_args.append(
                     _is_tuple_expr_of_literals(arg)

src/gt4py/next/iterator/transforms/global_tmps.py

@@ -93,18 +93,13 @@ def _transform_by_pattern(
             domain_expr = domain.as_expr()
 
             assert isinstance(tmp_expr.type, ts.TypeSpec)
-            tmp_names: str | tuple[str | tuple, ...] = type_info.apply_to_primitive_constituents(
-                lambda x: uids.sequential_id(),
-                tmp_expr.type,
-                tuple_constructor=lambda *elements: tuple(elements),
-            )
-            tmp_dtypes: ts.ScalarType | tuple[ts.ScalarType | tuple, ...] = (
-                type_info.apply_to_primitive_constituents(
-                    type_info.extract_dtype,
-                    tmp_expr.type,
-                    tuple_constructor=lambda *elements: tuple(elements),
-                )
-            )
+            tmp_names: str | tuple[str | tuple, ...] = type_info.type_tree_map(
+                result_collection_constructor=lambda elements: tuple(elements)
+            )(lambda x: uids.sequential_id())(tmp_expr.type)
+
+            tmp_dtypes: ts.ScalarType | tuple[ts.ScalarType | tuple, ...] = type_info.type_tree_map(
+                result_collection_constructor=lambda elements: tuple(elements)
+            )(type_info.extract_dtype)(tmp_expr.type)
 
             # allocate temporary for all tuple elements
             def allocate_temporary(tmp_name: str, dtype: ts.ScalarType):

src/gt4py/next/iterator/type_system/inference.py

@@ -478,9 +478,8 @@ def visit_Temporary(self, node: itir.Temporary, *, ctx) -> ts.FieldType | ts.Tup
         assert isinstance(domain, it_ts.DomainType)
         assert domain.dims != "unknown"
         assert node.dtype
-        return type_info.apply_to_primitive_constituents(
-            lambda dtype: ts.FieldType(dims=domain.dims, dtype=dtype),  # type: ignore[arg-type]  # ensured by domain.dims != "unknown" above
-            node.dtype,
+        return type_info.type_tree_map(lambda dtype: ts.FieldType(dims=domain.dims, dtype=dtype))(
+            node.dtype
         )
 
     def visit_IfStmt(self, node: itir.IfStmt, *, ctx) -> None:

src/gt4py/next/iterator/type_system/type_synthesizer.py

@@ -258,7 +258,7 @@ def _convert_as_fieldop_input_to_iterator(
         input_dims = []
 
     element_type: ts.DataType
-    element_type = type_info.apply_to_primitive_constituents(type_info.extract_dtype, input_)
+    element_type = type_info.type_tree_map(type_info.extract_dtype)(input_)
 
     # handle neighbor / sparse input fields
     defined_dims = []
@@ -311,12 +311,11 @@ def applied_as_fieldop(*fields) -> ts.FieldType | ts.DeferredType:
             offset_provider_type=offset_provider_type,
         )
         assert isinstance(stencil_return, ts.DataType)
-        return type_info.apply_to_primitive_constituents(
+        return type_info.type_tree_map(
             lambda el_type: ts.FieldType(dims=domain.dims, dtype=el_type)
             if domain.dims != "unknown"
-            else ts.DeferredType(constraint=ts.FieldType),
-            stencil_return,
-        )
+            else ts.DeferredType(constraint=ts.FieldType)
+        )(stencil_return)
 
     return applied_as_fieldop
 

src/gt4py/next/type_system/type_info.py

@@ -25,7 +25,7 @@
 import numpy as np
 
 from gt4py.eve.utils import XIterable, xiter
-from gt4py.next import common
+from gt4py.next import common, utils
 from gt4py.next.type_system import type_specifications as ts
 
 
@@ -197,6 +197,11 @@ def apply_to_primitive_constituents(
         return fun(*symbol_types)
 
 
+type_tree_map = utils.tree_map(
+    collection_type=ts.TupleType, result_collection_constructor=lambda x: ts.TupleType(types=[*x])
+)
+
+
 def extract_dtype(symbol_type: ts.TypeSpec) -> ts.ScalarType:
     """
     Extract the data type from ``symbol_type`` if it is either `FieldType` or `ScalarType`.

src/gt4py/next/utils.py

@@ -79,31 +79,36 @@ def tree_map(
 @overload
 def tree_map(
     *,
-    collection_type: type | tuple[type, ...] = tuple,
-    result_collection_constructor: Optional[type | Callable] = None,
-) -> Callable[
-    [Callable[_P, _R]], Callable[..., Any]
-]: ...  # TODO(havogt): if result_collection_constructor is Callable, improve typing
+    collection_type: type | tuple[type, ...],
+    result_collection_constructor: Optional[type] = None,
+) -> Callable[[Callable[_P, _R]], Callable[..., _R | tuple[_R | tuple, ...]]]: ...
 
 
 def tree_map(
-    fun: Optional[Callable[_P, _R]] = None,
-    *,
+    *args: Callable[_P, _R],
     collection_type: type | tuple[type, ...] = tuple,
     result_collection_constructor: Optional[type | Callable] = None,
-) -> Callable[..., _R | tuple[_R | tuple, ...]] | Callable[[Callable[_P, _R]], Callable[..., Any]]:
+) -> (
+    Callable[..., _R | tuple[_R | tuple, ...]]
+    | Callable[[Callable[_P, _R]], Callable[..., _R | tuple[_R | tuple, ...]]]
+    | _R
+    | tuple[_R | tuple, ...]
+):
     """
-    Apply `fun` to each entry of (possibly nested) collections (by default `tuple`s).
+    Apply `args[0]` to each entry of (possibly nested) collections (by default `tuple`s).
 
     Args:
-        fun: Function to apply to each entry of the collection.
+        args[0]: Function to apply to each entry of the collection.
         collection_type: Type of the collection to be traversed. Can be a single type or a tuple of types.
-        result_collection_constructor: Type of the collection to be returned. If `None` the same type as `collection_type` is used.
+        result_collection_constructor: Constructor of the collection to be returned. If `None` the same type as `collection_type` is used.
 
     Examples:
         >>> tree_map(lambda x: x + 1)(((1, 2), 3))
         ((2, 3), 4)
 
+        >>> tree_map(lambda x: x + 1, ((1, 2), 3))
+        ((2, 3), 4)
+
         >>> tree_map(lambda x, y: x + y)(((1, 2), 3), ((4, 5), 6))
         ((5, 7), 9)
 
@@ -129,7 +134,24 @@ def tree_map(
             )
         result_collection_constructor = collection_type
 
-    if fun:
+    if len(args) == 0:
+        return functools.partial(
+            tree_map,
+            collection_type=collection_type,
+            result_collection_constructor=result_collection_constructor,
+        )
+
+    if callable(args[0]):
+        fun = args[0]
+        colls = args[1:]
+
+        if len(colls) == 0:
+            return functools.partial(
+                tree_map,
+                fun,
+                collection_type=collection_type,
+                result_collection_constructor=result_collection_constructor,
+            )
 
         @functools.wraps(fun)
         def impl(*args: Any | tuple[Any | tuple, ...]) -> _R | tuple[_R | tuple, ...]:
@@ -140,14 +162,12 @@ def impl(*args: Any | tuple[Any | tuple, ...]) -> _R | tuple[_R | tuple, ...]:
                 assert result_collection_constructor is not None
                 return result_collection_constructor(impl(*arg) for arg in zip(*args))
 
-            return fun(  # type: ignore[call-arg, misc] # mypy not smart enough
+            return fun(  # type: ignore[call-arg] # mypy not smart enough
                 *cast(_P.args, args)
             )  # mypy doesn't understand that `args` at this point is of type `_P.args`
 
-        return impl
-    else:
-        return functools.partial(
-            tree_map,
-            collection_type=collection_type,
-            result_collection_constructor=result_collection_constructor,
-        )
+        return impl(*colls)
+
+    raise TypeError(
+        "tree_map() can be used as decorator with optional kwarg `collection_type` and `result_collection_constructor`, or with a function and collection."
+    )

tests/next_tests/unit_tests/test_utils.py

@@ -9,14 +9,45 @@
 import pytest
 
 from gt4py.next import utils
+from gt4py.next.type_system import type_specifications as ts
+from gt4py.next.type_system import type_info
+from gt4py.next.common import Field
+from numpy import int64
+
+
+def test_tree_map_scalar():
+    @utils.tree_map(collection_type=ts.ScalarType, result_collection_constructor=tuple)
+    def testee(x):
+        return x + 1
+
+    assert testee(1) == (2)
+
+
+def test_apply_to_primitive_constituents():
+    int_type = ts.ScalarType(kind=ts.ScalarKind.INT64)
+    tuple_type = ts.TupleType(types=[ts.TupleType(types=[int_type, int_type]), int_type])
+
+    tree = type_info.type_tree_map(
+        lambda primitive_type: ts.FieldType(dims=[], dtype=primitive_type)
+    )(tuple_type)
+
+    prim = type_info.apply_to_primitive_constituents(
+        lambda primitive_type: ts.FieldType(dims=[], dtype=primitive_type), tuple_type
+    )
+
+    assert tree == prim
 
 
 def test_tree_map_default():
+    expected_result = ((2, 3), 4)
+
     @utils.tree_map
     def testee(x):
         return x + 1
 
-    assert testee(((1, 2), 3)) == ((2, 3), 4)
+    assert testee(((1, 2), 3)) == expected_result
+    assert utils.tree_map(lambda x: x + 1)(((1, 2), 3)) == expected_result
+    assert utils.tree_map(lambda x: x + 1, ((1, 2), 3)) == expected_result
 
 
 def test_tree_map_multi_arg():