feat[next]: Add support for tuple comprehensions

Author: tehrengruber

docs/development/ADRs/next/0028-Homogeneous_Tuple_Comprehensions.md

@@ -0,0 +1,41 @@
+---
+tags: []
+---
+
+# Homogeneous Tuple Comprehensions
+
+- **Status**: valid
+- **Authors**: Till Ehrengruber (@tehrengruber), Sara Faghih-Naini (@SF-N)
+- **Created**: 2026-08-27
+- **Updated**: 2026-08-27
+
+In the context of tuple comprehensions in the field-view frontend, facing the constraint that every FOAST node carries exactly one type, we decided to support only homogeneous iterables — all elements of the same type — and reject heterogeneous ones in the type deduction.
+
+## Context
+
+Tuple comprehensions, e.g. `tuple(2.0 * el for el in (a, b))`, are typed and lowered with a single mapper: one target (`el`) and one element expression (`2.0 * el`), shared by all elements of the iterable. In FOAST every node has a single `type` attribute, so the target symbol — and consequently every node in the element expression — can only be typed once. If the iterable's elements had different types, the mapper would need a different type per element, i.e. per-element re-typing (monomorphization) of the element expression, which the FOAST type system does not support.
+
+The same constraint exists at the GTIR level: the ITIR type inference also stores a single type per node (and asserts on conflicting re-assignment), so the single `map_tuple` lambda used to lower comprehensions over variable-length tuples can only have one function type. For variable-length iterables heterogeneity cannot occur in the first place, since `VarArgType` describes all elements with a single element type. Rejecting heterogeneous iterables in the FOAST type deduction just surfaces the error earliest, with a source location.
+
+## Decision
+
+Only homogeneous iterables are supported, both fixed-length and variable-length. Heterogeneous ones are rejected in the type deduction. E.g., with `a`, `b`, `c`, `d` of equal type:
+
+```python
+tuple(2.0 * el for el in (a, b))  # supported
+tuple(2.0 * el for el in (a(V2E), b(V2E)))  # supported
+tuple(local_el + el for local_el, el in ((a(V2E), b), (c(V2E), d)))  # supported
+tuple(2.0 * el for el in (a(V2E), b))  # rejected: local vs. non-local element
+```
+
+Note that homogeneity applies to the iterable's elements as a whole: in the third example each element is a pair of a local and a non-local field, but all elements share that same tuple type, so each target symbol still has a single consistent type.
+
+## Consequences
+
+- Typing and lowering stay simple: the element expression is visited once, with one type per node.
+- Computations over differently-typed elements cannot be written as a comprehension; they must be spelled out per element.
+- The restriction could be lifted for fixed-length iterables by typing the mapper generically: the target symbol gets a type variable bounded by the valid element types, so a single type per node still suffices during type deduction. After `map_tuple` expansion the mapper is instantiated once per element, and each instance can then be specialized to its concrete element type. This is possible as a follow-up without breaking existing code, since it only widens the set of accepted programs. For variable-length iterables there is nothing to lift: heterogeneity cannot occur, as `VarArgType` has a single element type by construction.
+
+## References
+
+- PR [#2833](https://github.com/GridTools/gt4py/pull/2833) (tuple comprehension support)

docs/development/ADRs/next/README.md

@@ -26,6 +26,7 @@ Writing a new ADR is simple:
 - [0002 - Field View Lowering](0002-Field_View_Lowering.md)
 - [0010 - Domain in Field View](0010-Domain_in_Field_View.md)
 - [0013 - Scalar vs 0d-Fields](0013-Scalar_vs_0d_Fields.md)
+- [0028 - Homogeneous Tuple Comprehensions](0028-Homogeneous_Tuple_Comprehensions.md)
 
 ### Iterator IR #iterator
 

src/gt4py/next/ffront/dialect_parser.py

@@ -55,7 +55,10 @@
     ),
     ast.SetComp: ("set comprehension", ()),
     ast.DictComp: ("dictionary comprehension", ()),
-    ast.GeneratorExp: ("generator expression", ()),
+    ast.GeneratorExp: (
+        "generator expression",
+        ("Generator expressions are only supported as the argument of 'tuple(...)'.",),
+    ),
     ast.Lambda: (
         "'lambda' expression",
         ("Define a separate function decorated with '@field_operator' instead.",),

src/gt4py/next/ffront/field_operator_ast.py

@@ -12,6 +12,7 @@
 
 from gt4py import eve
 from gt4py.eve import Coerced, Node, SourceLocation, SymbolName, SymbolRef, datamodels
+from gt4py.eve.extended_typing import MaybeNestedInTuple
 from gt4py.eve.traits import SymbolTableTrait
 from gt4py.eve.type_definitions import StrEnum
 from gt4py.next import utils
@@ -98,6 +99,35 @@ class TupleExpr(Expr):
     elts: list[Expr]
 
 
+# TODO(tehrengruber): extend this to supported nested tuple comprehension.
+#  e.g. `tuple(element_expr for child in nested_tuple for grand_child in child)`
+#  would be represented by:
+#  ```
+#  class TupleComprehension(Expr):                              # ruff: noqa: ERA001
+#    inner: TupleComprehensionMapper | NestedTupleCompr         # ruff: noqa: ERA001
+#  class NestedTupleCompr(Expr, SymbolTableTrait):              # ruff: noqa: ERA001
+#    params: tuple[DataSymbol]                                  # ruff: noqa: ERA001
+#    body: TupleComprehension                                   # ruff: noqa: ERA001
+#  ```
+class TupleComprehension(Expr):
+    """
+    tuple(element_expr for target in iterable)
+    Note: The structure here differs from the one in the Python AST. Here we group target and
+    element expression in order to cleanly nest by the symbols being introduced, whereas in
+    the Python AST target and iterable are grouped into generator nodes.
+    """
+
+    inner: TupleComprehensionMapper
+    iterable: Expr
+
+
+# This is essentially a lambda. The difference is that for a lambda we might not know the type of
+# the args; therefore this is named differently at the moment.
+class TupleComprehensionMapper(LocatedNode, SymbolTableTrait):
+    target: MaybeNestedInTuple[DataSymbol]
+    element_expr: Expr
+
+
 class UnaryOp(Expr):
     op: dialect_ast_enums.UnaryOperator
     operand: Expr

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

@@ -12,6 +12,7 @@
 import gt4py.next.ffront.field_operator_ast as foast
 from gt4py import eve
 from gt4py.eve import NodeTranslator, NodeVisitor, traits
+from gt4py.eve.extended_typing import MaybeNestedInTuple
 from gt4py.next import common, errors
 from gt4py.next.common import Dimension, DimensionKind, promote_dims
 from gt4py.next.ffront import (
@@ -24,6 +25,7 @@
 from gt4py.next.ffront.foast_passes import utils as foast_utils
 from gt4py.next.iterator import builtins
 from gt4py.next.type_system import type_info, type_specifications as ts, type_translation
+from gt4py.next.utils import tree_map
 
 
 OperatorNodeT = TypeVar("OperatorNodeT", bound=foast.LocatedNode)
@@ -743,6 +745,117 @@ def visit_TupleExpr(self, node: foast.TupleExpr, **kwargs: Any) -> foast.TupleEx
         new_type = ts.TupleType(types=[element.type for element in new_elts])
         return foast.TupleExpr(elts=new_elts, type=new_type, location=node.location)
 
+    def _deduce_tuple_comprehension_target_type(
+        self,
+        target: MaybeNestedInTuple[foast.Symbol],
+        element_type: ts.TypeSpec,
+        **kwargs: Any,
+    ) -> MaybeNestedInTuple[foast.Symbol]:
+        """
+        Deduce the types of the loop target, e.g. `(a, b)` in `tuple(a + b for (a, b) in it)`.
+
+        Each target symbol starts out with a deferred type; revisiting it with `refine_type`
+        replaces that by the constituent of `element_type` at the symbol's unpacking position.
+        """
+
+        @tree_map(with_path_arg=True)
+        def process_target(target_el: foast.Symbol, path: tuple[int, ...]) -> foast.Symbol:
+            try:
+                type_ = element_type
+                for i in path:
+                    if not isinstance(type_, ts.TupleType) or len(type_.types) <= i:
+                        raise IndexError()
+                    type_ = type_.types[i]
+                return self.visit(target_el, refine_type=type_, **kwargs)
+            except IndexError:
+                raise errors.DSLError(
+                    target_el.location, f"Cannot unpack non-iterable '{type_}' object."
+                ) from None
+
+        return process_target(target)
+
+    def _deduce_tuple_comprehension_mapper(
+        self,
+        node: foast.TupleComprehension,
+        target: MaybeNestedInTuple[foast.Symbol],
+        element_type: ts.DataType,
+        **kwargs: Any,
+    ) -> foast.TupleComprehensionMapper:
+        """
+        Deduce the per-element part, e.g. `a + b for (a, b)` in `tuple(a + b for (a, b) in it)`.
+
+        Refines the target symbols against `element_type` and then types the element
+        expression with those symbols in scope (a fresh child symbol table).
+        """
+        inner_kwargs = {**kwargs, "symtable": kwargs["symtable"].new_child()}
+        new_target = self._deduce_tuple_comprehension_target_type(
+            target, element_type, **inner_kwargs
+        )
+        return foast.TupleComprehensionMapper(
+            target=new_target,
+            element_expr=self.visit(node.inner.element_expr, **inner_kwargs),
+            location=node.location,
+        )
+
+    def visit_TupleComprehension(
+        self, node: foast.TupleComprehension, **kwargs: Any
+    ) -> foast.TupleComprehension:
+        # The target symbols are visited in `_deduce_tuple_comprehension_target_type`,
+        # where their types are refined against the iterable's element type.
+        target = node.inner.target
+        iterable = self.visit(node.iterable, **kwargs)
+
+        if isinstance(iterable.type, ts.TupleType):
+            if len(iterable.type.types) == 0:
+                raise errors.DSLError(
+                    iterable.location,
+                    "Cannot iterate over an empty tuple in a tuple comprehension.",
+                )
+            if not all(
+                isinstance(element_type, ts.DataType) for element_type in iterable.type.types
+            ):
+                raise errors.DSLError(
+                    iterable.location,
+                    "Tuple comprehension iterable elements must be data types.",
+                )
+
+            element_types = cast(list[ts.DataType], iterable.type.types)
+            # Only homogeneous iterables are supported, see ADR 0028.
+            if not all(element_type == element_types[0] for element_type in element_types):
+                raise NotImplementedError(
+                    "Tuple comprehensions over fixed-length tuples require all iterable "
+                    "elements to have the same type (see ADR 0028)."
+                )
+            new_mapper = self._deduce_tuple_comprehension_mapper(
+                node, target, element_types[0], **kwargs
+            )
+            result = foast.TupleComprehension(
+                inner=new_mapper,
+                iterable=iterable,
+                location=node.location,
+                type=ts.TupleType(types=[new_mapper.element_expr.type for _ in element_types]),
+            )
+            return result
+        elif isinstance(iterable.type, ts.VarArgType):
+            element_type = iterable.type.element_type
+            new_mapper = self._deduce_tuple_comprehension_mapper(
+                node, target, element_type, **kwargs
+            )
+            element_expr = new_mapper.element_expr
+            return_type = ts.VarArgType(element_type=element_expr.type)
+
+            return foast.TupleComprehension(
+                inner=new_mapper,
+                iterable=iterable,
+                location=node.location,
+                type=return_type,
+            )
+        else:
+            raise errors.DSLError(
+                iterable.location,
+                f"Iterable in generator expression must be a tuple, got '{iterable.type}'.",
+            )
+
     def visit_Call(self, node: foast.Call, **kwargs: Any) -> foast.Call:
         new_func = self.visit(node.func, **kwargs)
         new_args = self.visit(node.args, **kwargs)

src/gt4py/next/ffront/foast_pretty_printer.py

@@ -120,6 +120,18 @@ def apply(cls, node: foast.LocatedNode, **kwargs: Any) -> str:  # type: ignore[o
 
     UnaryOp = as_fmt("{op}{operand}")
 
+    def visit_TupleComprehensionMapper(
+        self, node: foast.TupleComprehensionMapper, **kwargs: Any
+    ) -> str:
+        element_expr = self.visit(node.element_expr, **kwargs)
+        target = self.visit(node.target, **kwargs)
+        return f"{element_expr} for {target}"
+
+    def visit_TupleComprehension(self, node: foast.TupleComprehension, **kwargs: Any) -> str:
+        mapper = self.visit(node.inner, **kwargs)
+        iterable = self.visit(node.iterable, **kwargs)
+        return f"tuple(({mapper} in {iterable}))"
+
     def visit_UnaryOp(self, node: foast.UnaryOp, **kwargs: Any) -> str:
         if node.op is dialect_ast_enums.UnaryOperator.NOT:
             op = "not "

src/gt4py/next/ffront/foast_to_gtir.py

@@ -8,6 +8,7 @@
 
 
 import dataclasses
+import functools
 from typing import Any, Callable, Optional
 
 from gt4py import eve
@@ -259,6 +260,75 @@ def visit_Subscript(self, node: foast.Subscript, **kwargs: Any) -> itir.Expr:
     def visit_TupleExpr(self, node: foast.TupleExpr, **kwargs: Any) -> itir.Expr:
         return im.make_tuple(*[self.visit(el, **kwargs) for el in node.elts])
 
+    def _bind_tuple_comprehension_target(
+        self,
+        comprehension_target: itir.Sym | tuple,
+        element_expr: itir.Expr,
+        iterable_element: itir.Expr | str,
+    ) -> itir.Expr:
+        """
+        Wrap `element_expr` in a `let` binding the `comprehension_target` to `iterable_element`.
+
+        For `2.0 * a + b for (a, b) in iterable`:
+        - `comprehension_target`: `(a, b)`
+        - `element_expr`: `2.0 * a + b`
+        - `iterable_element`: the current element of `iterable`
+
+        returns
+        `let a = iterable_element[0], b = iterable_element[1] in element_expr`.
+        """
+        if isinstance(comprehension_target, itir.Sym):
+            return im.let(comprehension_target, iterable_element)(element_expr)
+
+        flat_targets = utils.flatten_nested_tuple(comprehension_target)
+        nested_target_values = utils.tree_map(
+            lambda _, path: functools.reduce(
+                lambda element, index: im.tuple_get(index, element), path, iterable_element
+            ),
+            with_path_arg=True,
+        )(comprehension_target)
+
+        flat_target_values = utils.flatten_nested_tuple(nested_target_values)  # type: ignore[arg-type]
+
+        target_bindings = tuple(zip(flat_targets, flat_target_values, strict=True))
+        return im.let(*target_bindings)(element_expr)  # type: ignore[arg-type]
+
+    def visit_TupleComprehension(self, node: foast.TupleComprehension, **kwargs: Any) -> itir.Expr:
+        # Only homogeneous iterables — all elements of the same type — are supported;
+        # heterogeneous ones are rejected in the type deduction (see ADR 0028).
+        comprehension_target = self.visit(node.inner.target, **kwargs)
+        element_expr = self.visit(node.inner.element_expr, **kwargs)
+        iterable_expr = self.visit(node.iterable, **kwargs)
+        iterable_type = node.iterable.type
+
+        lower_body_for_iterable_element = functools.partial(
+            self._bind_tuple_comprehension_target, comprehension_target, element_expr
+        )
+
+        if isinstance(iterable_type, ts.TupleType):
+            assert isinstance(node.type, ts.TupleType)
+            iterable_value_name = next(self.uid_generator["__tuple_comprh"])
+
+            fixed_tuple_elements = [
+                lower_body_for_iterable_element(im.tuple_get(element_index, iterable_value_name))
+                for element_index in range(len(iterable_type.types))
+            ]
+
+            result_tuple = im.make_tuple(*fixed_tuple_elements)
+            return im.let(iterable_value_name, iterable_expr)(result_tuple)
+
+        assert isinstance(iterable_type, ts.VarArgType)
+        assert isinstance(node.type, ts.VarArgType)
+        if isinstance(comprehension_target, itir.Sym):
+            map_tuple_lambda = im.lambda_(comprehension_target)(element_expr)
+        else:
+            iterable_element_param = next(self.uid_generator["__tuple_comprh"])
+            map_tuple_lambda = im.lambda_(iterable_element_param)(
+                lower_body_for_iterable_element(iterable_element_param)
+            )
+
+        return im.call(im.call("map_tuple")(map_tuple_lambda))(iterable_expr)
+
     def visit_UnaryOp(self, node: foast.UnaryOp, **kwargs: Any) -> itir.Expr:
         # TODO(tehrengruber): extend iterator ir to support unary operators
         dtype = type_info.extract_dtype(node.type)

src/gt4py/next/ffront/foast_to_past.py

@@ -113,9 +113,12 @@ def __call__(self, inp: ConcreteFOASTOperatorDef) -> ConcretePASTProgramDef:
             *partial_program_type.definition.kw_only_args.keys(),
         ]
         assert isinstance(type_, ts.CallableType)
-        assert arg_types[-1] == type_info.return_type(
+        return_type = type_info.return_type(
             type_, with_args=list(arg_types), with_kwargs=kwarg_types
         )
+        # Not equality: variadic comprehensions give a `VarArg[...]` return type,
+        # while 'out' is a concrete tuple.
+        assert type_info.is_concretizable(return_type, arg_types[-1])
         assert args_names[-1] == "out"
 
         params_decl: list[past.Symbol] = [