bug[next]: fix lowering of astype on tuples containing scalars

src/gt4py/next/ffront/foast_to_gtir.py

@@ -319,15 +319,19 @@ def _visit_astype(self, node: foast.Call, **kwargs: Any) -> itir.Expr:
         assert len(node.args) == 2 and isinstance(node.args[1], foast.Name)
         obj, new_type = self.visit(node.args[0], **kwargs), node.args[1].id
 
-        def create_cast(expr: itir.Expr) -> itir.FunCall:
-            return im.as_fieldop(
-                im.lambda_("__val")(im.call("cast_")(im.deref("__val"), str(new_type)))
-            )(expr)
+        def create_cast(expr: itir.Expr, t: ts.TypeSpec) -> itir.FunCall:
+            if isinstance(t, ts.FieldType):
+                return im.as_fieldop(
+                    im.lambda_("__val")(im.call("cast_")(im.deref("__val"), str(new_type)))
+                )(expr)
+            else:
+                assert isinstance(t, ts.ScalarType)
+                return im.call("cast_")(expr, str(new_type))
 
         if not isinstance(node.type, ts.TupleType):  # to keep the IR simpler
-            return create_cast(obj)
+            return create_cast(obj, node.type)
 
-        return lowering_utils.process_elements(create_cast, obj, node.type)
+        return lowering_utils.process_elements(create_cast, obj, node.type, with_type=True)
 
     def _visit_where(self, node: foast.Call, **kwargs: Any) -> itir.FunCall:
         if not isinstance(node.type, ts.TupleType):  # to keep the IR simpler

src/gt4py/next/ffront/lowering_utils.py

@@ -102,6 +102,7 @@ def process_elements(
     process_func: Callable[..., itir.Expr],
     objs: itir.Expr | Iterable[itir.Expr],
     current_el_type: ts.TypeSpec,
+    with_type: bool = False,
 ) -> itir.FunCall:
     """
     Recursively applies a processing function to all primitive constituents of a tuple.
@@ -112,13 +113,17 @@ def process_elements(
         objs: The object whose elements are to be transformed.
         current_el_type: A type with the same structure as the elements of `objs`. The leaf-types
             are not used and thus not relevant.
+        with_type: If True, the last argument passed to `process_func` will be its type.
     """
     if isinstance(objs, itir.Expr):
         objs = (objs,)
 
     let_ids = tuple(f"__val_{eve_utils.content_hash(obj)}" for obj in objs)
     body = _process_elements_impl(
-        process_func, tuple(im.ref(let_id) for let_id in let_ids), current_el_type
+        process_func,
+        tuple(im.ref(let_id) for let_id in let_ids),
+        current_el_type,
+        with_type=with_type,
     )
 
     return im.let(*(zip(let_ids, objs, strict=True)))(body)
@@ -131,6 +136,7 @@ def _process_elements_impl(
     process_func: Callable[..., itir.Expr],
     _current_el_exprs: Iterable[T],
     current_el_type: ts.TypeSpec,
+    with_type: bool,
 ) -> itir.Expr:
     if isinstance(current_el_type, ts.TupleType):
         result = im.make_tuple(
@@ -141,13 +147,17 @@ def _process_elements_impl(
                         im.tuple_get(i, current_el_expr) for current_el_expr in _current_el_exprs
                     ),
                     current_el_type.types[i],
+                    with_type=with_type,
                 )
                 for i in range(len(current_el_type.types))
             )
         )
     elif type_info.contains_local_field(current_el_type):
         raise NotImplementedError("Processing fields with local dimension is not implemented.")
     else:
-        result = process_func(*_current_el_exprs)
+        if with_type:
+            result = process_func(*_current_el_exprs, current_el_type)
+        else:
+            result = process_func(*_current_el_exprs)
 
     return result

tests/next_tests/unit_tests/ffront_tests/test_foast_to_gtir.py

@@ -19,6 +19,7 @@
 from gt4py.next import (
     astype,
     broadcast,
+    common,
     float32,
     float64,
     int32,
@@ -27,7 +28,6 @@
     min_over,
     neighbor_sum,
     where,
-    common,
 )
 from gt4py.next.ffront import type_specifications as ts_ffront
 from gt4py.next.ffront.ast_passes import single_static_assign as ssa
@@ -291,6 +291,18 @@ def foo(a: gtx.Field[[TDim], float64]):
     assert lowered.expr == reference
 
 
+def test_astype_scalar():
+    def foo(a: float64):
+        return astype(a, int32)
+
+    parsed = FieldOperatorParser.apply_to_function(foo)
+    lowered = FieldOperatorLowering.apply(parsed)
+
+    reference = im.call("cast_")("a", "int32")
+
+    assert lowered.expr == reference
+
+
 def test_astype_tuple():
     def foo(a: tuple[gtx.Field[[TDim], float64], gtx.Field[[TDim], float64]]):
         return astype(a, int32)
@@ -311,6 +323,24 @@ def foo(a: tuple[gtx.Field[[TDim], float64], gtx.Field[[TDim], float64]]):
     assert lowered_inlined.expr == reference
 
 
+def test_astype_tuple_scalar_and_field():
+    def foo(a: tuple[gtx.Field[[TDim], float64], float64]):
+        return astype(a, int32)
+
+    parsed = FieldOperatorParser.apply_to_function(foo)
+    lowered = FieldOperatorLowering.apply(parsed)
+    lowered_inlined = inline_lambdas.InlineLambdas.apply(lowered)
+
+    reference = im.make_tuple(
+        im.as_fieldop(im.lambda_("__val")(im.call("cast_")(im.deref("__val"), "int32")))(
+            im.tuple_get(0, "a")
+        ),
+        im.call("cast_")(im.tuple_get(1, "a"), "int32"),
+    )
+
+    assert lowered_inlined.expr == reference
+
+
 def test_astype_nested_tuple():
     def foo(
         a: tuple[