feat[next]: Elementwise operations on Tuples

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/__init__.py

@@ -31,6 +31,7 @@
     Field,
     GridType,
     UnitRange,
+    XTuple,
     as_non_staggered,
     domain,
     flip_staggered,
@@ -118,6 +119,7 @@
     "DimensionKind",
     "Dims",
     "Field",
+    "XTuple",
     "CartesianConnectivity",
     "Connectivity",
     "GridType",

src/gt4py/next/common.py

@@ -57,6 +57,87 @@
 class Dims(tuple[Unpack[ShapeTs]]): ...
 
 
+class XTuple(tuple[Unpack[ShapeTs]]):
+    """
+    Tuple on which binary arithmetic and logical operators apply element-wise.
+
+    Non-tuple operands (e.g. scalars, fields) are broadcast against the tuple structure.
+    Nested element-wise behavior requires nested `XTuple`s; plain `tuple` operands are
+    rejected, mirroring the type deduction rules of the DSL.
+    """
+
+    def _elementwise_op(self, other: Any, op: Callable[[Any, Any], Any]) -> XTuple:
+        self_elems: tuple[Any, ...] = tuple(self)
+        if isinstance(other, XTuple):
+            other_elems: tuple[Any, ...] = tuple(other)
+            if len(self_elems) != len(other_elems):
+                raise ValueError(
+                    f"Element-wise operations require 'XTuple's of equal length, "
+                    f"got {len(self_elems)} and {len(other_elems)}."
+                )
+            return XTuple(op(el, other_el) for el, other_el in zip(self_elems, other_elems))
+        if isinstance(other, tuple):
+            raise TypeError("Element-wise operations require 'XTuple' operands, got 'tuple'.")
+        return XTuple(op(el, other) for el in self_elems)
+
+    def __add__(self, other: Any) -> XTuple:
+        return self._elementwise_op(other, lambda a, b: a + b)
+
+    def __radd__(self, other: Any) -> XTuple:
+        return self._elementwise_op(other, lambda a, b: b + a)
+
+    def __sub__(self, other: Any) -> XTuple:
+        return self._elementwise_op(other, lambda a, b: a - b)
+
+    def __rsub__(self, other: Any) -> XTuple:
+        return self._elementwise_op(other, lambda a, b: b - a)
+
+    def __mul__(self, other: Any) -> XTuple:
+        return self._elementwise_op(other, lambda a, b: a * b)
+
+    def __rmul__(self, other: Any) -> XTuple:
+        return self._elementwise_op(other, lambda a, b: b * a)
+
+    def __truediv__(self, other: Any) -> XTuple:
+        return self._elementwise_op(other, lambda a, b: a / b)
+
+    def __rtruediv__(self, other: Any) -> XTuple:
+        return self._elementwise_op(other, lambda a, b: b / a)
+
+    def __floordiv__(self, other: Any) -> XTuple:
+        return self._elementwise_op(other, lambda a, b: a // b)
+
+    def __rfloordiv__(self, other: Any) -> XTuple:
+        return self._elementwise_op(other, lambda a, b: b // a)
+
+    def __mod__(self, other: Any) -> XTuple:
+        return self._elementwise_op(other, lambda a, b: a % b)
+
+    def __rmod__(self, other: Any) -> XTuple:
+        return self._elementwise_op(other, lambda a, b: b % a)
+
+    def __pow__(self, other: Any) -> XTuple:
+        return self._elementwise_op(other, lambda a, b: a**b)
+
+    def __and__(self, other: Any) -> XTuple:
+        return self._elementwise_op(other, lambda a, b: a & b)
+
+    def __rand__(self, other: Any) -> XTuple:
+        return self._elementwise_op(other, lambda a, b: b & a)
+
+    def __or__(self, other: Any) -> XTuple:
+        return self._elementwise_op(other, lambda a, b: a | b)
+
+    def __ror__(self, other: Any) -> XTuple:
+        return self._elementwise_op(other, lambda a, b: b | a)
+
+    def __xor__(self, other: Any) -> XTuple:
+        return self._elementwise_op(other, lambda a, b: a ^ b)
+
+    def __rxor__(self, other: Any) -> XTuple:
+        return self._elementwise_op(other, lambda a, b: b ^ a)
+
+
 DimsT = TypeVar("DimsT", bound=Dims, covariant=True)
 
 Tag: TypeAlias = str

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