From e5b676acbb663b05bad412d8db90a4a46bc977f2 Mon Sep 17 00:00:00 2001 From: Till Ehrengruber Date: Thu, 27 Aug 2026 09:48:26 +0000 Subject: [PATCH 01/24] fix[next]: 'concat_where' crashes on local (neighbor-list) fields The GTIR 'concat_where' type synthesizer derives the result dims via 'type_info.promote(tb, fb)', which asserted that the promoted dtypes are 'ScalarType'. A local field carries a 'ListType' dtype at the GTIR level, so any 'concat_where' with a local-field branch failed with a bare 'AssertionError' during type inference -- and it did so even though both branches have the identical dtype, since the assertion is on the dtype's kind, not on the operands differing. Let 'promote' handle 'ListType' the same way it handles 'ScalarType': both promote only between equal types, so the two cases collapse into a single check. The docstring records that a 'ListType' only ever reaches 'promote' from the ITIR level, because the frontend represents the same concept as a field with a local dimension in 'dims' and a scalar dtype. The combination was previously untested ('test_concat_where.py' had no local-dimension coverage), so the latent assert never fired in CI. --- .../ffront/foast_passes/type_deduction.py | 8 ++- src/gt4py/next/type_system/type_info.py | 34 ++++++--- .../ffront_tests/test_concat_where.py | 72 ++++++++++++++++++- 3 files changed, 100 insertions(+), 14 deletions(-) diff --git a/src/gt4py/next/ffront/foast_passes/type_deduction.py b/src/gt4py/next/ffront/foast_passes/type_deduction.py index c9f51ad080..6d25c3edb3 100644 --- a/src/gt4py/next/ffront/foast_passes/type_deduction.py +++ b/src/gt4py/next/ffront/foast_passes/type_deduction.py @@ -871,8 +871,12 @@ def _visit_math_built_in(self, node: foast.Call, **kwargs: Any) -> foast.Call: ) elif func_name in fbuiltins.BINARY_MATH_NUMBER_BUILTIN_NAMES: try: - return_type = type_info.promote( - *((cast(ts.FieldType | ts.ScalarType, arg.type)) for arg in node.args) + return_type = cast( + # a `ListType` only occurs at the ITIR level, never in the frontend + ts.FieldType | ts.ScalarType, + type_info.promote( + *((cast(ts.FieldType | ts.ScalarType, arg.type)) for arg in node.args) + ), ) except ValueError as ex: raise errors.DSLError(node.location, error_msg_preamble) from ex diff --git a/src/gt4py/next/type_system/type_info.py b/src/gt4py/next/type_system/type_info.py index ac8467a5f6..1dd404fde3 100644 --- a/src/gt4py/next/type_system/type_info.py +++ b/src/gt4py/next/type_system/type_info.py @@ -9,7 +9,7 @@ import functools import types from collections.abc import Callable, Collection, Iterable, Iterator -from typing import Any, Final, Literal, Sequence, Type, TypeGuard, TypeVar, cast, overload +from typing import Any, Final, Literal, Sequence, Type, TypeGuard, TypeVar, overload import numpy as np @@ -556,15 +556,20 @@ def is_concretizable(symbol_type: ts.TypeSpec, to_type: ts.TypeSpec) -> bool: def promote( - *types: ts.FieldType | ts.ScalarType, always_field: bool = False -) -> ts.FieldType | ts.ScalarType: + *types: ts.FieldType | ts.ScalarType | ts.ListType, always_field: bool = False +) -> ts.FieldType | ts.ScalarType | ts.ListType: """ - Promote a set of field or scalar types to a common type. + Promote a set of field, scalar or list types to a common type. The resulting type is defined on all dimensions of the arguments, respecting the individual order of the dimensions of each argument (see :func:`common.promote_dims` for more details). + A `ListType` is the dtype of a local (neighbor-list) field and only ever reaches this + function from the ITIR level. The frontend represents the same concept as a field with + a local dimension in `dims` and a scalar dtype, so it never passes a list here. Lists + promote exactly like scalars, i.e. only between equal types. + >>> dtype = ts.ScalarType(kind=ts.ScalarKind.INT64) >>> I, J, K = (common.Dimension(value=dim) for dim in ["I", "J", "K"]) >>> promoted: ts.FieldType = promote( @@ -578,18 +583,27 @@ def promote( ... ) >>> promoted.dims == [I, J, K] and promoted.dtype == dtype True + + >>> V2E = common.Dimension(value="V2E", kind=common.DimensionKind.LOCAL) + >>> list_dtype = ts.ListType(element_type=dtype, offset_type=V2E) + >>> promote( + ... ts.FieldType(dims=[I], dtype=list_dtype), + ... ts.FieldType(dims=[I, J], dtype=list_dtype), + ... ) == ts.FieldType(dims=[I, J], dtype=list_dtype) + True """ - if not always_field and all(isinstance(type_, ts.ScalarType) for type_ in types): + # Lists are only reached from the ITIR level (see above) and behave like scalars here: + # both promote only between equal types. + if not always_field and all(isinstance(type_, (ts.ScalarType, ts.ListType)) for type_ in types): if not all(type_ == types[0] for type_ in types): - raise ValueError("Could not promote scalars of different dtype (not implemented).") - if not all(type_.shape is None for type_ in types): # type: ignore[union-attr] + raise ValueError("Could not promote dtypes of different type (not implemented).") + if isinstance(types[0], ts.ScalarType) and types[0].shape is not None: raise NotImplementedError("Shape promotion not implemented.") return types[0] elif all(isinstance(type_, (ts.ScalarType, ts.FieldType)) for type_ in types): dims = common.promote_dims(*(extract_dims(type_) for type_ in types)) - extracted_dtypes = [extract_dtype(type_) for type_ in types] - assert all(isinstance(dtype, ts.ScalarType) for dtype in extracted_dtypes) - dtype = cast(ts.ScalarType, promote(*extracted_dtypes)) # type: ignore[arg-type] # checked is `ScalarType` + dtype = promote(*(extract_dtype(type_) for type_ in types)) + assert isinstance(dtype, (ts.ScalarType, ts.ListType)) return ts.FieldType(dims=dims, dtype=dtype) raise TypeError("Expected a 'FieldType' or 'ScalarType'.") diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_concat_where.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_concat_where.py index 577f5a520e..f46c166e98 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_concat_where.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_concat_where.py @@ -8,13 +8,23 @@ import numpy as np import pytest -from next_tests.integration_tests.cases import IDim, JDim, KDim, cartesian_case +from next_tests.integration_tests.cases import ( + E2V, + E2VDim, + Edge, + IDim, + JDim, + KDim, + cartesian_case, + unstructured_case, +) from gt4py import next as gtx -from gt4py.next import broadcast +from gt4py.next import broadcast, common, neighbor_sum from gt4py.next.ffront.experimental import concat_where from next_tests.integration_tests import cases from next_tests.integration_tests.cases_utils import ( exec_alloc_descriptor, + mesh_descriptor, ) pytestmark = pytest.mark.uses_concat_where @@ -390,6 +400,64 @@ def ref(interior0, boundary0, interior1, boundary1): cases.verify_with_default_data(cartesian_case, testee, ref) +@pytest.mark.uses_unstructured_shift +@pytest.mark.uses_sparse_fields +def test_with_local_field(unstructured_case, static_domains: bool): + @gtx.field_operator(static_domains=static_domains) + def testee(a: cases.VField, b: cases.VField) -> cases.EField: + t = concat_where(Edge < 2, a(E2V), b(E2V)) + return neighbor_sum(t, axis=E2VDim) + + e2v_table = unstructured_case.offset_provider["E2V"].asnumpy() + edge_mask = np.arange(unstructured_case.default_sizes[Edge]) < 2 + cases.verify_with_default_data( + unstructured_case, + testee, + ref=lambda a, b: np.sum( + np.where(edge_mask[:, np.newaxis], a[e2v_table], b[e2v_table]), + axis=1, + initial=0, + where=e2v_table != common._DEFAULT_SKIP_VALUE, + ), + ) + + +@pytest.mark.uses_tuple_returns +@pytest.mark.uses_unstructured_shift +@pytest.mark.uses_sparse_fields +def test_with_tuples_of_local_fields(unstructured_case, static_domains: bool): + @gtx.field_operator(static_domains=static_domains) + def testee( + a: cases.VField, + b: cases.VField, + c: cases.VField, + d: cases.VField, + ) -> tuple[cases.EField, cases.EField]: + t = concat_where(Edge < 2, (a(E2V), b(E2V)), (c(E2V), d(E2V))) + return neighbor_sum(t[0], axis=E2VDim), neighbor_sum(t[1], axis=E2VDim) + + e2v_table = unstructured_case.offset_provider["E2V"].asnumpy() + edge_mask = np.arange(unstructured_case.default_sizes[Edge]) < 2 + cases.verify_with_default_data( + unstructured_case, + testee, + ref=lambda a, b, c, d: ( + np.sum( + np.where(edge_mask[:, np.newaxis], a[e2v_table], c[e2v_table]), + axis=1, + initial=0, + where=e2v_table != common._DEFAULT_SKIP_VALUE, + ), + np.sum( + np.where(edge_mask[:, np.newaxis], b[e2v_table], d[e2v_table]), + axis=1, + initial=0, + where=e2v_table != common._DEFAULT_SKIP_VALUE, + ), + ), + ) + + def test_nested_conditions_with_empty_branches(cartesian_case, static_domains: bool): @gtx.field_operator(static_domains=static_domains) def testee(interior: cases.IField, boundary: cases.IField, N: gtx.int32) -> cases.IField: From 8643e6d67e8d3b63198bf7a0a2c082724be99d90 Mon Sep 17 00:00:00 2001 From: Till Ehrengruber Date: Thu, 27 Aug 2026 19:58:03 +0000 Subject: [PATCH 02/24] Add 'concat_where' test with nested tuple branches (moved from #2833) --- .../ffront_tests/test_concat_where.py | 45 +++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_concat_where.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_concat_where.py index f46c166e98..6ed7b244c4 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_concat_where.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_concat_where.py @@ -400,6 +400,51 @@ def ref(interior0, boundary0, interior1, boundary1): cases.verify_with_default_data(cartesian_case, testee, ref) +@pytest.mark.uses_tuple_returns +def test_with_nested_tuples(cartesian_case, static_domains: bool): + @gtx.field_operator(static_domains=static_domains) + def testee( + interior0: cases.IJKField, + boundary0: cases.IJField, + interior1: cases.IJKField, + boundary1: cases.IJField, + interior2: cases.IJKField, + boundary2: cases.IJField, + ) -> tuple[cases.IJKField, tuple[cases.IJKField, cases.IJKField]]: + return concat_where( + KDim == 0, + (boundary0, (boundary1, boundary2)), + (interior0, (interior1, interior2)), + ) + + interiors = tuple(cases.allocate(cartesian_case, testee, f"interior{i}")() for i in range(3)) + boundaries = tuple(cases.allocate(cartesian_case, testee, f"boundary{i}")() for i in range(3)) + out = cases.allocate(cartesian_case, testee, cases.RETURN)() + + k = np.arange(0, cartesian_case.default_sizes[KDim]) + refs = tuple( + np.where( + k[np.newaxis, np.newaxis, :] == 0, + boundary.asnumpy()[:, :, np.newaxis], + interior.asnumpy(), + ) + for boundary, interior in zip(boundaries, interiors) + ) + + cases.verify( + cartesian_case, + testee, + interiors[0], + boundaries[0], + interiors[1], + boundaries[1], + interiors[2], + boundaries[2], + out=out, + ref=(refs[0], (refs[1], refs[2])), + ) + + @pytest.mark.uses_unstructured_shift @pytest.mark.uses_sparse_fields def test_with_local_field(unstructured_case, static_domains: bool): From 7b21e9c870660b8035870b5e8770cc71bda3f407 Mon Sep 17 00:00:00 2001 From: Till Ehrengruber Date: Thu, 27 Aug 2026 20:09:16 +0000 Subject: [PATCH 03/24] feat[next]: Add support for tuple comprehensions --- .../0028-Homogeneous_Tuple_Comprehensions.md | 41 ++++ docs/development/ADRs/next/README.md | 1 + src/gt4py/next/ffront/dialect_parser.py | 5 +- src/gt4py/next/ffront/field_operator_ast.py | 30 +++ .../ffront/foast_passes/type_deduction.py | 113 +++++++++++ src/gt4py/next/ffront/foast_pretty_printer.py | 12 ++ src/gt4py/next/ffront/foast_to_gtir.py | 70 +++++++ src/gt4py/next/ffront/foast_to_past.py | 5 +- src/gt4py/next/ffront/func_to_foast.py | 80 ++++++-- src/gt4py/next/ffront/lowering_utils.py | 22 ++- .../next/ffront/past_passes/type_deduction.py | 4 +- src/gt4py/next/type_system/type_info.py | 8 + .../next/type_system/type_specifications.py | 9 + .../next/type_system/type_translation.py | 53 ++++-- tests/next_tests/integration_tests/cases.py | 11 ++ .../feature_tests/ffront_tests/test_tuples.py | 99 +++++++++- .../feature_tests/ffront_tests/test_where.py | 70 ++++++- .../ffront_tests/test_diagnostic_messages.py | 35 ++++ .../ffront_tests/test_foast_to_gtir.py | 179 +++++++++++++++++- .../ffront_tests/test_func_to_foast.py | 33 ++++ .../test_type_translation.py | 20 +- 21 files changed, 840 insertions(+), 60 deletions(-) create mode 100644 docs/development/ADRs/next/0028-Homogeneous_Tuple_Comprehensions.md diff --git a/docs/development/ADRs/next/0028-Homogeneous_Tuple_Comprehensions.md b/docs/development/ADRs/next/0028-Homogeneous_Tuple_Comprehensions.md new file mode 100644 index 0000000000..9b3b34e4f3 --- /dev/null +++ b/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) diff --git a/docs/development/ADRs/next/README.md b/docs/development/ADRs/next/README.md index 1bf9d21812..45b6fd9240 100644 --- a/docs/development/ADRs/next/README.md +++ b/docs/development/ADRs/next/README.md @@ -29,6 +29,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 diff --git a/src/gt4py/next/ffront/dialect_parser.py b/src/gt4py/next/ffront/dialect_parser.py index a7c27a3181..c464ffd4dc 100644 --- a/src/gt4py/next/ffront/dialect_parser.py +++ b/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.",), diff --git a/src/gt4py/next/ffront/field_operator_ast.py b/src/gt4py/next/ffront/field_operator_ast.py index f84dd53d62..f95f9f285a 100644 --- a/src/gt4py/next/ffront/field_operator_ast.py +++ b/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 diff --git a/src/gt4py/next/ffront/foast_passes/type_deduction.py b/src/gt4py/next/ffront/foast_passes/type_deduction.py index 6d25c3edb3..9e52486921 100644 --- a/src/gt4py/next/ffront/foast_passes/type_deduction.py +++ b/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) diff --git a/src/gt4py/next/ffront/foast_pretty_printer.py b/src/gt4py/next/ffront/foast_pretty_printer.py index 8b2e369501..79826983a3 100644 --- a/src/gt4py/next/ffront/foast_pretty_printer.py +++ b/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 " diff --git a/src/gt4py/next/ffront/foast_to_gtir.py b/src/gt4py/next/ffront/foast_to_gtir.py index 480a05812f..283253bce7 100644 --- a/src/gt4py/next/ffront/foast_to_gtir.py +++ b/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) diff --git a/src/gt4py/next/ffront/foast_to_past.py b/src/gt4py/next/ffront/foast_to_past.py index 9a560b7ff8..1bbd3d9991 100644 --- a/src/gt4py/next/ffront/foast_to_past.py +++ b/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] = [ diff --git a/src/gt4py/next/ffront/func_to_foast.py b/src/gt4py/next/ffront/func_to_foast.py index 14dceb25d1..d2926a90fc 100644 --- a/src/gt4py/next/ffront/func_to_foast.py +++ b/src/gt4py/next/ffront/func_to_foast.py @@ -11,9 +11,9 @@ import ast import textwrap import typing -from typing import Any, Type import gt4py.eve as eve +from gt4py.eve.extended_typing import Any, MaybeNestedInTuple from gt4py.next import errors from gt4py.next.ffront import ( dialect_ast_enums, @@ -324,7 +324,7 @@ def visit_Assign( if not isinstance(target, ast.Name): raise errors.DSLError(self.get_location(node), "Can only assign to names.") new_value = self.visit(node.value) - constraint_type: Type[ts.DataType] = ts.DataType + constraint_type: type[ts.DataType] = ts.DataType if isinstance(new_value, foast.TupleExpr): constraint_type = ts.TupleType elif ( @@ -542,24 +542,76 @@ def visit_NotEq(self, node: ast.NotEq, **kwargs: Any) -> foast.CompareOperator: return foast.CompareOperator.NOTEQ def _verify_builtin_type_constructor(self, node: ast.Call) -> None: - if len(node.args) > 0: - arg = node.args[0] - if not ( - isinstance(arg, ast.Constant) - or (isinstance(arg, ast.UnaryOp) and isinstance(arg.operand, ast.Constant)) - ): - raise errors.DSLError( - self.get_location(node), - f"'{self._func_name(node)}()' only takes literal arguments.", - ) + assert isinstance(node.func, ast.Name) + (arg,) = node.args + if not ( + isinstance(arg, ast.Constant) + or (isinstance(arg, ast.UnaryOp) and isinstance(arg.operand, ast.Constant)) + or (node.func.id == "tuple" and isinstance(arg, ast.GeneratorExp)) + ): + allowed = ( + "literal arguments or a generator expression" + if node.func.id == "tuple" + else "literal arguments" + ) + raise errors.DSLError( + self.get_location(node), + f"'{self._func_name(node)}()' only takes {allowed}.", + ) def _func_name(self, node: ast.Call) -> str: return node.func.id # type: ignore[attr-defined] # We want this to fail if the attribute does not exist unexpectedly. - def visit_Call(self, node: ast.Call, **kwargs: Any) -> foast.Call: - # TODO(tehrengruber): is this still needed or redundant with the checks in type deduction? + def visit_Call(self, node: ast.Call, **kwargs: Any) -> foast.Call | foast.TupleComprehension: if isinstance(node.func, ast.Name): func_name = self._func_name(node) + + if ( + func_name == "tuple" + and len(node.args) == 1 + and isinstance(gen_expr := node.args[0], ast.GeneratorExp) + ): + if len(gen_expr.generators) != 1: + raise errors.DSLError( + self.get_location(gen_expr), + "Nested generator expressions are not supported.", + hints=["Use a single 'for' clause iterating over one tuple."], + ) + if gen_expr.generators[0].ifs != []: + raise errors.DSLError( + self.get_location(gen_expr.generators[0].ifs[0]), + "Conditionals are not supported in generator expressions.", + notes=[ + ( + "The length of the resulting tuple must be known at compile time, " + "but an 'if' filter makes it depend on runtime values." + ) + ], + ) + + def parse_target(target: ast.expr) -> MaybeNestedInTuple[foast.DataSymbol]: + if isinstance(target, ast.Tuple): + return tuple(parse_target(el) for el in target.elts) + assert isinstance(target, ast.Name) and isinstance(target.ctx, ast.Store) + return foast.DataSymbol( + id=target.id, + location=self.get_location(target), + type=ts.DeferredType(constraint=None), + ) + + target = parse_target(gen_expr.generators[0].target) + + return foast.TupleComprehension( + inner=foast.TupleComprehensionMapper( + target=target, + element_expr=self.visit(gen_expr.elt, **kwargs), + location=self.get_location(node), + ), + iterable=self.visit(gen_expr.generators[0].iter, **kwargs), + location=self.get_location(node), + ) + + # TODO(tehrengruber): is this still needed or redundant with the checks in type deduction? if func_name in fbuiltins.TYPE_BUILTIN_NAMES: self._verify_builtin_type_constructor(node) diff --git a/src/gt4py/next/ffront/lowering_utils.py b/src/gt4py/next/ffront/lowering_utils.py index c4e35c9e18..f7a5d26991 100644 --- a/src/gt4py/next/ffront/lowering_utils.py +++ b/src/gt4py/next/ffront/lowering_utils.py @@ -20,6 +20,7 @@ def process_elements( objs: itir.Expr | Iterable[itir.Expr], current_el_type: ts.TypeSpec, arg_types: Optional[Iterable[ts.TypeSpec]] = None, + with_path_arg: bool = False, ) -> itir.FunCall: """ Recursively applies a processing function to all primitive constituents of a tuple or @@ -34,6 +35,7 @@ def process_elements( arg_types: If provided, a tuple of the type of each argument is passed to `process_func` as last argument. Note, that `arg_types` might coincide with `(current_el_type,)*len(objs)`, but not necessarily, in case of implicit broadcasts. + with_path_arg: If true, the index path to the current leaf is passed as last argument. """ if isinstance(objs, itir.Expr): objs = (objs,) @@ -47,6 +49,8 @@ def process_elements( tuple(im.ref(var_name) for var_name in var_names), current_el_type, arg_types=arg_types, + path=(), + with_path_arg=with_path_arg, ) return im.let(*bound_vars.items())(body) @@ -60,6 +64,8 @@ def _process_elements_impl( _current_el_exprs: Iterable[T], current_el_type: ts.TypeSpec, arg_types: Optional[Iterable[ts.TypeSpec]], + path: tuple[int, ...], + with_path_arg: bool, ) -> itir.Expr: if isinstance(current_el_type, (ts.TupleType, ts.NamedCollectionType)): result = im.make_tuple( @@ -70,17 +76,23 @@ def _process_elements_impl( im.tuple_get(i, current_el_expr) for current_el_expr in _current_el_exprs ), current_el_type.types[i], - arg_types=tuple(arg_t.types[i] for arg_t in arg_types) # type: ignore[attr-defined] # guaranteed by the requirement that `current_el_type` and each element of `arg_types` have the same tuple structure - if arg_types is not None - else None, + arg_types=( + tuple(arg_t.types[i] for arg_t in arg_types) # type: ignore[attr-defined] # guaranteed by the requirement that `current_el_type` and each element of `arg_types` have the same tuple structure + if arg_types is not None + else None + ), + path=(*path, i), + with_path_arg=with_path_arg, ) for i in range(len(current_el_type.types)) ) ) else: if arg_types is not None: - result = process_func(*_current_el_exprs, arg_types) + result = process_func( + *_current_el_exprs, arg_types, *((path,) if with_path_arg else ()) + ) else: - result = process_func(*_current_el_exprs) + result = process_func(*_current_el_exprs, *((path,) if with_path_arg else ())) return result diff --git a/src/gt4py/next/ffront/past_passes/type_deduction.py b/src/gt4py/next/ffront/past_passes/type_deduction.py index 9d021ceb51..fcc75b7374 100644 --- a/src/gt4py/next/ffront/past_passes/type_deduction.py +++ b/src/gt4py/next/ffront/past_passes/type_deduction.py @@ -248,7 +248,9 @@ def visit_Call(self, node: past.Call, **kwargs: Any) -> past.Call: operator_return_type = type_info.return_type( new_func.type, with_args=arg_types, with_kwargs=kwarg_types ) - if operator_return_type != new_kwargs["out"].type: + # Not equality: variadic comprehensions give a `VarArg[...]` return type, + # while 'out' is a concrete tuple. + if not type_info.is_compatible_type(operator_return_type, new_kwargs["out"].type): raise ValueError( "Expected keyword argument 'out' to be of " f"type '{operator_return_type}', got " diff --git a/src/gt4py/next/type_system/type_info.py b/src/gt4py/next/type_system/type_info.py index 1dd404fde3..e3651b1c62 100644 --- a/src/gt4py/next/type_system/type_info.py +++ b/src/gt4py/next/type_system/type_info.py @@ -550,6 +550,14 @@ def is_concretizable(symbol_type: ts.TypeSpec, to_type: ts.TypeSpec) -> bool: or issubclass(type_class(to_type), symbol_type.constraint) ): return True + if isinstance(symbol_type, ts.VarArgType) and isinstance(to_type, ts.VarArgType): + return is_concretizable(symbol_type.element_type, to_type.element_type) + if isinstance(symbol_type, ts.VarArgType) and isinstance(to_type, ts.TupleType): + if len(to_type.types) == 0 or ( + all(type_ == to_type.types[0] for type_ in to_type.types) + and is_concretizable(symbol_type.element_type, to_type.types[0]) + ): + return True elif is_concrete(symbol_type): return symbol_type == to_type return False diff --git a/src/gt4py/next/type_system/type_specifications.py b/src/gt4py/next/type_system/type_specifications.py index 59ac40f0f3..1016eba34e 100644 --- a/src/gt4py/next/type_system/type_specifications.py +++ b/src/gt4py/next/type_system/type_specifications.py @@ -148,6 +148,15 @@ def __len__(self) -> int: return len(self.types) +class VarArgType(DataType): + """Represents a variable number of arguments of the same type.""" + + element_type: DataType + + def __str__(self) -> str: + return f"VarArg[{self.element_type}]" + + class AnyPythonType: """Marker type representing any Python type which cannot be used for instantiation. diff --git a/src/gt4py/next/type_system/type_translation.py b/src/gt4py/next/type_system/type_translation.py index 850013ddfa..4ff4d3970b 100644 --- a/src/gt4py/next/type_system/type_translation.py +++ b/src/gt4py/next/type_system/type_translation.py @@ -75,6 +75,16 @@ def make_constructor_type(type_spec: ts.TypeSpec) -> ts.ConstructorType: ) ) + case ts.DeferredType(constraint=ts.TupleType): + return ts.ConstructorType( + definition=ts.FunctionType( + pos_only_args=[ts.DeferredType(constraint=None)], + pos_or_kw_args={}, + kw_only_args={}, + returns=ts.DeferredType(constraint=ts.VarArgType), + ) + ) + case ts.NamedCollectionType() as named_tuple_type: type_ = pkgutil.resolve_name(named_tuple_type.original_python_type) pos_or_kw_args = {k: t for k, t in zip(type_spec.keys, type_spec.types)} @@ -147,7 +157,7 @@ def canonicalize_type_hint( *, globalns: Optional[dict[str, Any]] = None, localns: Optional[dict[str, Any]] = None, -) -> tuple[Any, tuple[Any, ...], tuple[Any, ...]]: +) -> tuple[Any, tuple[Any, ...] | None, tuple[Any, ...]]: """ Canonicalize python type annotations as a tuple of (canonical_type, type_args, annotated_extra_args). """ @@ -178,7 +188,8 @@ def canonicalize_type_hint( type_hint = _resolve_type_alias(type_hint) canonical_type = typing.get_origin(type_hint) or type_hint - args = typing.get_args(type_hint) + # In order to distinguish `tuple` from `tuple[()]`, the former returns None here. + args = typing.get_args(type_hint) if typing.get_origin(type_hint) else None return canonical_type, args, tuple(extra_args) @@ -199,17 +210,37 @@ def from_type_hint( match canonical_type: case builtins.tuple: - if not args: - raise ValueError(f"Tuple annotation '{type_hint}' requires at least one argument.") - if Ellipsis in args: - raise ValueError(f"Unbound tuples '{type_hint}' are not allowed.") - tuple_types = [from_type_hint_same_ns(arg) for arg in args] - assert all(isinstance(elem, ts.DataType) for elem in tuple_types) - return ts.TupleType(types=tuple_types) + # Fixed-length tuple, e.g. `tuple[int32, float64]` + if ( + isinstance(args, tuple) + and len(args) > 0 + and not any(arg is Ellipsis for arg in args) + ): + tuple_types = [from_type_hint_same_ns(arg) for arg in args] + assert all(isinstance(elem, ts.DataType) for elem in tuple_types) + return ts.TupleType(types=tuple_types) + # Variable-length tuple, e.g. `tuple[int32, ...]` + elif isinstance(args, tuple) and len(args) == 2 and args[1] is Ellipsis: + return ts.VarArgType(element_type=from_type_hint_same_ns(args[0])) + # Unparametrized annotation, i.e. `tuple` or `tuple[()]` + elif args is None or (isinstance(args, tuple) and len(args) == 0): + # TODO(tehrengruber): We use `DeferredType` until we have an actual representation + # for a generic type. + return ts.DeferredType(constraint=ts.TupleType) + else: + raise ValueError( + f"Tuple annotation '{type_hint}' must either " + f"be a list of concrete arguments (e.g. 'tuple[int]'), " + f"be a variadic tuple (e.g. 'tuple[int, ...]'), " + f"or have no arguments (e.g. 'tuple')." + ) case common.Field: - if (n_args := len(args)) != 2: - raise ValueError(f"Field type requires two arguments, got {n_args}: '{type_hint}'.") + # `args` is `None` for an unparametrized `Field` annotation + if args is None or len(args) != 2: + raise ValueError( + f"Field type requires two arguments, got {len(args or ())}: '{type_hint}'." + ) dims: list[common.Dimension] = [] dim_arg, dtype_arg = args dim_arg = ( diff --git a/tests/next_tests/integration_tests/cases.py b/tests/next_tests/integration_tests/cases.py index 08fb817856..508125c341 100644 --- a/tests/next_tests/integration_tests/cases.py +++ b/tests/next_tests/integration_tests/cases.py @@ -616,6 +616,15 @@ def _allocate_from_type( for t in types ) ) + case ts.VarArgType(element_type=element_type): + return tuple( + ( + _allocate_from_type( + case=case, arg_type=t, domain=domain, dtype=dtype, strategy=strategy + ) + for t in [element_type] * 3 # TODO: revisit + ) + ) case ts.NamedCollectionType(types=types) as named_collection_type_spec: container_constructor = ( named_collections.make_named_collection_constructor_from_type_spec( @@ -661,6 +670,8 @@ def get_param_size(param_type: ts.TypeSpec, sizes: dict[gtx.Dimension, int]) -> return sum([get_param_size(t, sizes=sizes) for t in types]) case ts.NamedCollectionType(types=types): return sum([get_param_size(t, sizes=sizes) for t in types]) + case ts.VarArgType(element_type=element_type): + return get_param_size(ts.TupleType(types=[element_type] * 3), sizes) # TODO: revisit case _: raise TypeError(f"Can not get size for parameter of type '{param_type}'.") diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py index 88a5a18af1..090f6a7e7e 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py @@ -10,14 +10,7 @@ import pytest import gt4py.next as gtx -from gt4py.next import ( - broadcast, - errors, - float64, - int32, - neighbor_sum, - utils as gt_utils, -) +from gt4py.next import broadcast, errors, float64, int32, neighbor_sum, utils as gt_utils from next_tests.integration_tests import cases from next_tests.integration_tests.cases import ( @@ -120,6 +113,94 @@ def testee(a: tuple[cases.IField, cases.IJField]) -> cases.IJField: ) +@pytest.mark.uses_tuple_args +def test_fixed_len_tuple_comprehension(cartesian_case): + @gtx.field_operator + def testee( + tracers: tuple[cases.IField, cases.IField], factor: int32 + ) -> tuple[cases.IField, cases.IField]: + return tuple(tracer * factor for tracer in tracers) + + cases.verify_with_default_data( + cartesian_case, + testee, + ref=lambda t, f: tuple(el * f for el in t), + ) + + +@pytest.mark.uses_tuple_args +def test_var_len_tuple_comprehension(cartesian_case): + @gtx.field_operator + def testee(tracers: tuple[cases.IField, ...], factor: int32) -> tuple[cases.IField, ...]: + return tuple(tracer * factor for tracer in tracers) + + cases.verify_with_default_data( + cartesian_case, + testee, + ref=lambda t, f: tuple(el * f for el in t), + ) + + +@pytest.mark.uses_tuple_args +def test_tuple_comprehension_other_fo(cartesian_case): + @gtx.field_operator + def inner(tracer: cases.IField, factor: int32) -> cases.IField: + return tracer * factor + + @gtx.field_operator + def testee(tracers: tuple[cases.IField, ...], factor: int32) -> tuple[cases.IField, ...]: + return tuple(inner(tracer, factor) for tracer in tracers) + + cases.verify_with_default_data( + cartesian_case, + testee, + ref=lambda t, f: tuple(el * f for el in t), + ) + + +@pytest.mark.uses_tuple_args +def test_nested_tuple_comprehension(cartesian_case): + @gtx.field_operator + def testee( + vals: tuple[tuple[cases.IField, ...], ...], factor: int32 + ) -> tuple[tuple[cases.IField, ...], ...]: + return tuple(tuple(grand_child * factor for grand_child in child) for child in vals) + + cases.verify_with_default_data( + cartesian_case, + testee, + ref=lambda t, f: tuple(tuple(grand_child * f for grand_child in child) for child in t), + ) + + +@pytest.mark.uses_tuple_args +def test_nested_tuple_comprehension_shadowing_names(cartesian_case): + @gtx.field_operator + def testee( + vals: tuple[tuple[cases.IField, ...], ...], factor: int32 + ) -> tuple[tuple[cases.IField, ...], ...]: + return tuple(tuple(child * factor for child in child) for child in vals) + + cases.verify_with_default_data( + cartesian_case, + testee, + ref=lambda t, f: tuple(tuple(child * f for child in child) for child in t), + ) + + +@pytest.mark.uses_tuple_args +def test_multi_target_tuple_comprehension(cartesian_case): + @gtx.field_operator + def testee(nested_tuple: tuple[tuple[int32, cases.IField], ...]) -> tuple[cases.IField, ...]: + return tuple(factor * tracer for factor, tracer in nested_tuple) + + cases.verify_with_default_data( + cartesian_case, + testee, + ref=lambda t: tuple(f * el for f, el in t), + ) + + @pytest.mark.uses_tuple_args @pytest.mark.uses_unstructured_shift @pytest.mark.xfail(reason="Iterator of tuple approach in lowering does not allow this.") @@ -254,5 +335,5 @@ def test_tuple_unpacking_too_few_values(cartesian_case): @gtx.field_operator(backend=cartesian_case.backend) def _invalid_unpack() -> tuple[int32, float64, int32]: - a, b, c = 1 + a, _b, _c = 1 return a diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_where.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_where.py index 7ea11a7b69..7b93d7892e 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_where.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_where.py @@ -7,16 +7,27 @@ # SPDX-License-Identifier: BSD-3-Clause import functools + import numpy as np -from typing import Tuple import pytest -from next_tests.integration_tests.cases import IDim, JDim, KDim, cartesian_case + from gt4py import next as gtx -from gt4py.next import float64, int32 -from gt4py.next.ffront.fbuiltins import where, broadcast +from gt4py.next import common, float64, int32, neighbor_sum +from gt4py.next.ffront.fbuiltins import broadcast, where + from next_tests.integration_tests import cases +from next_tests.integration_tests.cases import ( + E2V, + E2VDim, + IDim, + JDim, + KDim, + cartesian_case, + unstructured_case, +) from next_tests.integration_tests.cases_utils import ( exec_alloc_descriptor, + mesh_descriptor, ) @@ -113,6 +124,41 @@ def testee( ) +@pytest.mark.uses_tuple_returns +@pytest.mark.uses_unstructured_shift +def test_with_tuples_and_local_condition(unstructured_case): + @gtx.field_operator + def testee( + a: cases.VField, + b: cases.VField, + c: cases.VField, + d: cases.VField, + ) -> tuple[cases.EField, cases.EField]: + cond = a(E2V) > c(E2V) + t = where(cond, (a(E2V), b(E2V)), (c(E2V), d(E2V))) + return neighbor_sum(t[0], axis=E2VDim), neighbor_sum(t[1], axis=E2VDim) + + e2v_table = unstructured_case.offset_provider["E2V"].asnumpy() + cases.verify_with_default_data( + unstructured_case, + testee, + ref=lambda a, b, c, d: ( + np.sum( + np.where(a[e2v_table] > c[e2v_table], a[e2v_table], c[e2v_table]), + axis=1, + initial=0, + where=e2v_table != common._DEFAULT_SKIP_VALUE, + ), + np.sum( + np.where(a[e2v_table] > c[e2v_table], b[e2v_table], d[e2v_table]), + axis=1, + initial=0, + where=e2v_table != common._DEFAULT_SKIP_VALUE, + ), + ), + ) + + @pytest.mark.uses_tuple_returns def test_conditional_nested_tuple(cartesian_case): @gtx.field_operator @@ -124,7 +170,9 @@ def conditional_nested_tuple( return where(mask, ((a, b), (b, a)), ((5.0, 7.0), (7.0, 5.0))) size = cartesian_case.default_sizes[IDim] - mask = cartesian_case.as_field([IDim], np.random.choice(a=[False, True], size=size)) + mask = cartesian_case.as_field( + [IDim], np.random.default_rng().choice(a=[False, True], size=size) + ) a = cases.allocate(cartesian_case, conditional_nested_tuple, "a")() b = cases.allocate(cartesian_case, conditional_nested_tuple, "b")() @@ -158,7 +206,9 @@ def conditional( return where(mask, a, b) size = cartesian_case.default_sizes[IDim] - mask = cartesian_case.as_field([IDim], np.random.choice(a=[False, True], size=(size))) + mask = cartesian_case.as_field( + [IDim], np.random.default_rng().choice(a=[False, True], size=size) + ) a = cases.allocate(cartesian_case, conditional, "a")() b = cases.allocate(cartesian_case, conditional, "b")() out = cases.allocate(cartesian_case, conditional, cases.RETURN)() @@ -180,7 +230,9 @@ def conditional_promotion(mask: cases.IBoolField, a: cases.IFloatField) -> cases return where(mask, a, 10.0) size = cartesian_case.default_sizes[IDim] - mask = cartesian_case.as_field([IDim], np.random.choice(a=[False, True], size=(size))) + mask = cartesian_case.as_field( + [IDim], np.random.default_rng().choice(a=[False, True], size=size) + ) a = cases.allocate(cartesian_case, conditional_promotion, "a")() out = cases.allocate(cartesian_case, conditional_promotion, cases.RETURN)() ref = np.where(mask.asnumpy(), a.asnumpy(), 10.0) @@ -214,7 +266,9 @@ def conditional_program( conditional_shifted(mask, a, b, out=out) size = cartesian_case.default_sizes[IDim] + 1 - mask = cartesian_case.as_field([IDim], np.random.choice(a=[False, True], size=(size))) + mask = cartesian_case.as_field( + [IDim], np.random.default_rng().choice(a=[False, True], size=size) + ) a = cases.allocate(cartesian_case, conditional_program, "a").extend({IDim: (0, 1)})() b = cases.allocate(cartesian_case, conditional_program, "b").extend({IDim: (0, 1)})() out = cases.allocate(cartesian_case, conditional_shifted, cases.RETURN)() diff --git a/tests/next_tests/unit_tests/ffront_tests/test_diagnostic_messages.py b/tests/next_tests/unit_tests/ffront_tests/test_diagnostic_messages.py index 62072cab18..50f03b0886 100644 --- a/tests/next_tests/unit_tests/ffront_tests/test_diagnostic_messages.py +++ b/tests/next_tests/unit_tests/ffront_tests/test_diagnostic_messages.py @@ -356,3 +356,38 @@ def broken(a: BrokenFieldAlias) -> gtx.Field[[IDim], float64]: assert err.location.line == err.location.end_line assert err.location.end_column - err.location.column == len("a: BrokenFieldAlias") assert re.search(r"\| +\^{19}(?!\^)", str(err)), str(err) + + +def test_bare_generator_expression_points_to_tuple_constructor(): + def with_genexp(a: gtx.Field[[IDim], float64]) -> gtx.Field[[IDim], float64]: + b = (x for x in (a, a)) # noqa: F841 [unused-variable] + return a + + err = parse_error(with_genexp) + + assert isinstance(err, errors.UnsupportedPythonFeatureError) + assert err.message == "Unsupported Python syntax: generator expression." + assert any("'tuple(...)'" in hint for hint in err.hints) + + +def test_nested_generator_expression_suggests_single_for_clause(): + def with_nested_genexp(a: gtx.Field[[IDim], float64]) -> gtx.Field[[IDim], float64]: + b = tuple(x for x in (a, a) for y in (a, a)) # noqa: F841 [unused-variable] + return a + + err = parse_error(with_nested_genexp) + + assert err.message == "Nested generator expressions are not supported." + assert any("single 'for' clause" in hint for hint in err.hints) + + +def test_conditional_in_generator_expression_explains_static_length(): + def with_filtered_genexp(a: gtx.Field[[IDim], float64]) -> gtx.Field[[IDim], float64]: + b = tuple(x for x in (a, a) if True) # noqa: F841 [unused-variable] + return a + + err = parse_error(with_filtered_genexp) + + assert err.message == "Conditionals are not supported in generator expressions." + assert any("known at compile time" in note for note in err.notes) + assert "if True" in str(err) diff --git a/tests/next_tests/unit_tests/ffront_tests/test_foast_to_gtir.py b/tests/next_tests/unit_tests/ffront_tests/test_foast_to_gtir.py index 407088e4f8..9c4dd08266 100644 --- a/tests/next_tests/unit_tests/ffront_tests/test_foast_to_gtir.py +++ b/tests/next_tests/unit_tests/ffront_tests/test_foast_to_gtir.py @@ -42,7 +42,6 @@ from gt4py.next.type_system import type_specifications as ts, type_translation from gt4py.next.iterator import ir as itir - Edge = gtx.Dimension("Edge") Vertex = gtx.Dimension("Vertex") V2EDim = gtx.Dimension("V2E", gtx.DimensionKind.LOCAL) @@ -382,6 +381,184 @@ def foo( assert lowered_inlined.expr == reference +def test_fixed_len_tuple_comprehension(): + def foo(a: tuple[gtx.Field[[TDim], float64], gtx.Field[[TDim], float64]], factor: float64): + return tuple(el * factor for el in a) + + parsed = FieldOperatorParser.apply_to_function(foo) + lowered = FieldOperatorLowering.apply(parsed) + lowered_inlined = inline_lambdas.InlineLambdas.apply(lowered) + + iterable = "__tuple_comprh_0" + reference = im.let(iterable, "a")( + im.make_tuple( + im.let("el", im.tuple_get(0, iterable))(im.op_as_fieldop("multiplies")("el", "factor")), + im.let("el", im.tuple_get(1, iterable))(im.op_as_fieldop("multiplies")("el", "factor")), + ) + ) + reference_inlined = inline_lambdas.InlineLambdas.apply(reference) + + assert lowered_inlined.expr == reference_inlined + + +def test_var_len_tuple_comprehension_scalar(): + def foo(a: tuple[float64, ...]): + return tuple(2.0 * el for el in a) + + parsed = FieldOperatorParser.apply_to_function(foo) + lowered = FieldOperatorLowering.apply(parsed) + + reference = im.call( + im.call("map_tuple")(im.lambda_("el")(im.multiplies_(im.literal("2.0", "float64"), "el"))) + )("a") + + assert lowered.expr == reference + + +def test_var_len_tuple_comprehension_field(): + def foo(a: tuple[gtx.Field[[TDim], float64], ...], factor: float64): + return tuple(el * factor for el in a) + + parsed = FieldOperatorParser.apply_to_function(foo) + lowered = FieldOperatorLowering.apply(parsed) + + reference = im.call( + im.call("map_tuple")(im.lambda_("el")(im.op_as_fieldop("multiplies")("el", "factor"))) + )("a") + + assert lowered.expr == reference + + +def test_var_len_tuple_comprehension_field_with_local_dim(): + def foo(a: tuple[gtx.Field[[Vertex, V2EDim], float64], ...]): + return tuple(2.0 * el for el in a) + + parsed = FieldOperatorParser.apply_to_function(foo) + lowered = FieldOperatorLowering.apply(parsed) + + two = im.literal("2.0", "float64") + reference = im.call( + im.call("map_tuple")( + im.lambda_("el")( + im.op_as_fieldop(im.map_list("multiplies"))( + im.op_as_fieldop("make_const_list")(two), "el" + ) + ) + ) + )("a") + + assert lowered.expr == reference + + +def test_fixed_len_tuple_comprehension_local_field(): + def foo(a: gtx.Field[[Edge], float64]): + return tuple(2.0 * el for el in (a(V2E), a(V2E))) + + parsed = FieldOperatorParser.apply_to_function(foo) + lowered = FieldOperatorLowering.apply(parsed) + lowered_inlined = inline_lambdas.InlineLambdas.apply(lowered) + + iterable = "__tuple_comprh_0" + two = im.literal("2.0", "float64") + reference = im.let( + iterable, + im.make_tuple(im.as_fieldop_neighbors("V2E", "a"), im.as_fieldop_neighbors("V2E", "a")), + )( + im.make_tuple( + im.let("el", im.tuple_get(0, iterable))( + im.op_as_fieldop(im.map_list("multiplies"))( + im.op_as_fieldop("make_const_list")(two), "el" + ) + ), + im.let("el", im.tuple_get(1, iterable))( + im.op_as_fieldop(im.map_list("multiplies"))( + im.op_as_fieldop("make_const_list")(two), "el" + ) + ), + ) + ) + reference_inlined = inline_lambdas.InlineLambdas.apply(reference) + + assert lowered_inlined.expr == reference_inlined + + +def test_fixed_len_tuple_comprehension_mixed_local_field_tuple_target(): + def foo( + a: gtx.Field[[Edge], float64], + b: gtx.Field[[Vertex], float64], + c: gtx.Field[[Edge], float64], + d: gtx.Field[[Vertex], float64], + ): + return tuple( + 2.0 * local_el + scalar_el for local_el, scalar_el in ((a(V2E), b), (c(V2E), d)) + ) + + parsed = FieldOperatorParser.apply_to_function(foo) + lowered = FieldOperatorLowering.apply(parsed) + lowered_inlined = inline_lambdas.InlineLambdas.apply(lowered) + + iterable = "__tuple_comprh_0" + two = im.literal("2.0", "float64") + local_term = im.op_as_fieldop(im.map_list("multiplies"))( + im.op_as_fieldop("make_const_list")(two), "local_el" + ) + element_expr = im.op_as_fieldop(im.map_list("plus"))( + local_term, im.op_as_fieldop("make_const_list")("scalar_el") + ) + tuple_el_0 = im.tuple_get(0, iterable) + tuple_el_1 = im.tuple_get(1, iterable) + mapped_tuple1 = im.make_tuple(im.as_fieldop_neighbors("V2E", "a"), "b") + mapped_tuple2 = im.make_tuple(im.as_fieldop_neighbors("V2E", "c"), "d") + reference = im.let(iterable, im.make_tuple(mapped_tuple1, mapped_tuple2))( + im.make_tuple( + im.let( + ("local_el", im.tuple_get(0, tuple_el_0)), + ("scalar_el", im.tuple_get(1, tuple_el_0)), + )(element_expr), + im.let( + ("local_el", im.tuple_get(0, tuple_el_1)), + ("scalar_el", im.tuple_get(1, tuple_el_1)), + )(element_expr), + ) + ) + reference_inlined = inline_lambdas.InlineLambdas.apply(reference) + + assert lowered_inlined.expr == reference_inlined + + +def test_fixed_len_tuple_comprehension_mixed_local_field(): + def foo(a: gtx.Field[[Edge], float64], b: gtx.Field[[Vertex], float64]): + return tuple(2.0 * el for el in (a(V2E), b)) + + with pytest.raises( + NotImplementedError, + match="fixed-length tuples require all iterable elements to have the same type", + ): + FieldOperatorParser.apply_to_function(foo) + + +def test_fixed_len_tuple_comprehension_mixed_field_domains(): + def foo(a: gtx.Field[[Edge], float64], b: gtx.Field[[Vertex], float64]): + return tuple(2.0 * el for el in (a, b)) + + with pytest.raises( + NotImplementedError, + match="fixed-length tuples require all iterable elements to have the same type", + ): + FieldOperatorParser.apply_to_function(foo) + + +def test_fixed_len_tuple_comprehension_tuple_target_mixed_element_types(): + def foo(a: gtx.Field[[Edge], float64], b: gtx.Field[[Vertex], float64]): + return tuple(2.0 * left + right for left, right in ((a(V2E), b), (b, b))) + + with pytest.raises( + NotImplementedError, + match="fixed-length tuples require all iterable elements to have the same type", + ): + FieldOperatorParser.apply_to_function(foo) + + def test_unary_minus(): def foo(inp: gtx.Field[[TDim], float64]): return -inp diff --git a/tests/next_tests/unit_tests/ffront_tests/test_func_to_foast.py b/tests/next_tests/unit_tests/ffront_tests/test_func_to_foast.py index a75e7bb031..bed07f0685 100644 --- a/tests/next_tests/unit_tests/ffront_tests/test_func_to_foast.py +++ b/tests/next_tests/unit_tests/ffront_tests/test_func_to_foast.py @@ -490,3 +490,36 @@ def tuple_index_failure( with pytest.raises(errors.DSLError, match=r"need .* literal"): _ = FieldOperatorParser.apply_to_function(tuple_index_failure) + + +def test_tuple_compr_non_tuple_iterable_failure(): + def testee(arg: float): + return tuple(_ for _ in arg) + + with pytest.raises( + errors.DSLError, + match=re.escape("Iterable in generator expression must be a tuple, got 'float64'."), + ): + _ = FieldOperatorParser.apply_to_function(testee) + + +def test_nested_tuple_compr_failure(): + def testee(nested_tuple: tuple[tuple[gtx.Field[[TDim], float64], ...], ...], factor: int32): + return tuple(grandchild * factor for child in nested_tuple for grandchild in child) + + with pytest.raises( + errors.DSLError, + match=re.escape("Nested generator expressions are not supported."), + ): + _ = FieldOperatorParser.apply_to_function(testee) + + +def test_tuple_compr_unpacking_failure(): + def testee(arg: tuple[int32, ...]): + return tuple(a * b for a, b in arg) + + with pytest.raises( + errors.DSLError, + match=re.escape("Cannot unpack non-iterable 'int32' object."), + ): + _ = FieldOperatorParser.apply_to_function(testee) diff --git a/tests/next_tests/unit_tests/type_system_tests/test_type_translation.py b/tests/next_tests/unit_tests/type_system_tests/test_type_translation.py index 91d644dcac..fe259f18d0 100644 --- a/tests/next_tests/unit_tests/type_system_tests/test_type_translation.py +++ b/tests/next_tests/unit_tests/type_system_tests/test_type_translation.py @@ -216,15 +216,17 @@ def test_invalid_symbol_types(): type_translation.from_type_hint("foo") # Tuples - with pytest.raises(ValueError, match="least one argument"): - type_translation.from_type_hint(typing.Tuple) - with pytest.raises(ValueError, match="least one argument"): - type_translation.from_type_hint(tuple) - - with pytest.raises(ValueError, match="Unbound tuples"): - type_translation.from_type_hint(tuple[int, ...]) - with pytest.raises(ValueError, match="Unbound tuples"): - type_translation.from_type_hint(typing.Tuple["float", ...]) + # Bare `tuple` and `typing.Tuple` (unparameterized) both return a DeferredType. + assert type_translation.from_type_hint(typing.Tuple) == ts.DeferredType(constraint=ts.TupleType) + assert type_translation.from_type_hint(tuple) == ts.DeferredType(constraint=ts.TupleType) + + # Variadic tuples (`tuple[T, ...]`) are now valid — returns a VarArgType. + assert type_translation.from_type_hint(tuple[int, ...]) == ts.VarArgType( + element_type=ts.ScalarType(kind=ts.ScalarKind.INT64) + ) + assert type_translation.from_type_hint(typing.Tuple["float", ...]) == ts.VarArgType( + element_type=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) + ) # Fields with pytest.raises(ValueError, match="Field type requires two arguments"): From 4a117afc563f4e6f0905a4178b521754a6ea2daa Mon Sep 17 00:00:00 2001 From: Till Ehrengruber Date: Thu, 27 Aug 2026 22:47:48 +0000 Subject: [PATCH 04/24] Address Fable review: builtin-constructor arity diagnostic and length-mismatch xfail test --- src/gt4py/next/ffront/func_to_foast.py | 5 ++++ .../feature_tests/ffront_tests/test_tuples.py | 25 +++++++++++++++++++ .../ffront_tests/test_func_to_foast.py | 14 +++++++++++ 3 files changed, 44 insertions(+) diff --git a/src/gt4py/next/ffront/func_to_foast.py b/src/gt4py/next/ffront/func_to_foast.py index d2926a90fc..b210027d5d 100644 --- a/src/gt4py/next/ffront/func_to_foast.py +++ b/src/gt4py/next/ffront/func_to_foast.py @@ -543,6 +543,11 @@ def visit_NotEq(self, node: ast.NotEq, **kwargs: Any) -> foast.CompareOperator: def _verify_builtin_type_constructor(self, node: ast.Call) -> None: assert isinstance(node.func, ast.Name) + if len(node.args) != 1: + raise errors.DSLError( + self.get_location(node), + f"'{self._func_name(node)}()' takes exactly one argument, got {len(node.args)}.", + ) (arg,) = node.args if not ( isinstance(arg, ast.Constant) diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py index 090f6a7e7e..392b900841 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py @@ -141,6 +141,31 @@ def testee(tracers: tuple[cases.IField, ...], factor: int32) -> tuple[cases.IFie ) +@pytest.mark.uses_tuple_args +@pytest.mark.xfail( + strict=True, + reason="The consistency of the lengths of variable-length tuple arguments and the " + "'out' argument is not checked anywhere; the mismatch only surfaces as an internal " + "error (e.g. a bare 'AssertionError' from 'utils.tree_map') deep in the toolchain.", +) +def test_var_len_tuple_length_mismatch_rejected(cartesian_case): + # The return length equals the length of `a` or `b` depending on a runtime value, so + # with `len(a) != len(b)` no single length can be consistent with `out` — this should + # be rejected with a proper diagnostic when the concrete argument types are known. + @gtx.field_operator + def testee( + a: tuple[cases.IField, ...], b: tuple[cases.IField, ...], c: bool + ) -> tuple[cases.IField, ...]: + return a if c else b + + a = cases.allocate(cartesian_case, testee, "a")()[:2] + b = cases.allocate(cartesian_case, testee, "b")() # allocated length: 3 + out = cases.allocate(cartesian_case, testee, cases.RETURN)()[:2] + + with pytest.raises(errors.DSLError): + cases.run(cartesian_case, testee, a, b, False, out=out) + + @pytest.mark.uses_tuple_args def test_tuple_comprehension_other_fo(cartesian_case): @gtx.field_operator diff --git a/tests/next_tests/unit_tests/ffront_tests/test_func_to_foast.py b/tests/next_tests/unit_tests/ffront_tests/test_func_to_foast.py index bed07f0685..b72e567a12 100644 --- a/tests/next_tests/unit_tests/ffront_tests/test_func_to_foast.py +++ b/tests/next_tests/unit_tests/ffront_tests/test_func_to_foast.py @@ -269,6 +269,20 @@ def cast_scalar_temp(): _ = FieldOperatorParser.apply_to_function(cast_scalar_temp) +def test_scalar_cast_wrong_arity(): + def cast_scalar_two_args(): + return int32(1, 2) + + with pytest.raises(errors.DSLError, match=r"takes exactly one argument, got 2"): + _ = FieldOperatorParser.apply_to_function(cast_scalar_two_args) + + def cast_scalar_no_args(): + return int32() + + with pytest.raises(errors.DSLError, match=r"takes exactly one argument, got 0"): + _ = FieldOperatorParser.apply_to_function(cast_scalar_no_args) + + def test_conditional_wrong_mask_type(): def conditional_wrong_mask_type(a: gtx.Field[[TDim], float64]) -> gtx.Field[[TDim], float64]: return where(a, a, a) From 666c8b4df8e687d0bb0a14025a66ce5074b2e0b8 Mon Sep 17 00:00:00 2001 From: Till Ehrengruber Date: Tue, 1 Sep 2026 07:34:02 +0000 Subject: [PATCH 05/24] Address review comments (unpack diagnostics, DSLError for heterogeneity, tuple[()] handling, empty-vararg test) --- .../ffront/foast_passes/type_deduction.py | 32 +++++++++++-------- src/gt4py/next/type_system/type_info.py | 2 ++ .../next/type_system/type_translation.py | 15 ++++++--- .../ffront_tests/test_foast_to_gtir.py | 12 +++---- .../ffront_tests/test_type_deduction.py | 18 +++++++++++ .../type_system_tests/test_type_info.py | 12 +++++++ .../test_type_translation.py | 6 ++++ 7 files changed, 73 insertions(+), 24 deletions(-) diff --git a/src/gt4py/next/ffront/foast_passes/type_deduction.py b/src/gt4py/next/ffront/foast_passes/type_deduction.py index 9e52486921..59215b8ef2 100644 --- a/src/gt4py/next/ffront/foast_passes/type_deduction.py +++ b/src/gt4py/next/ffront/foast_passes/type_deduction.py @@ -760,17 +760,20 @@ def _deduce_tuple_comprehension_target_type( @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 + type_ = element_type + for i in path: + if not isinstance(type_, ts.TupleType): + raise errors.DSLError( + target_el.location, f"Cannot unpack non-iterable '{type_}' object." + ) + if len(type_.types) <= i: + raise errors.DSLError( + target_el.location, + f"Not enough values to unpack (expected at least {i + 1}, " + f"got {len(type_.types)}).", + ) + type_ = type_.types[i] + return self.visit(target_el, refine_type=type_, **kwargs) return process_target(target) @@ -822,9 +825,10 @@ def visit_TupleComprehension( 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)." + raise errors.DSLError( + iterable.location, + "Tuple comprehensions over fixed-length tuples with differently typed " + "iterable elements are not implemented (see ADR 0028).", ) new_mapper = self._deduce_tuple_comprehension_mapper( node, target, element_types[0], **kwargs diff --git a/src/gt4py/next/type_system/type_info.py b/src/gt4py/next/type_system/type_info.py index e3651b1c62..f4f5288770 100644 --- a/src/gt4py/next/type_system/type_info.py +++ b/src/gt4py/next/type_system/type_info.py @@ -553,6 +553,8 @@ def is_concretizable(symbol_type: ts.TypeSpec, to_type: ts.TypeSpec) -> bool: if isinstance(symbol_type, ts.VarArgType) and isinstance(to_type, ts.VarArgType): return is_concretizable(symbol_type.element_type, to_type.element_type) if isinstance(symbol_type, ts.VarArgType) and isinstance(to_type, ts.TupleType): + # The empty tuple is a valid concretization of any variable-length tuple: there + # are no elements whose type could conflict with the element type. if len(to_type.types) == 0 or ( all(type_ == to_type.types[0] for type_ in to_type.types) and is_concretizable(symbol_type.element_type, to_type.types[0]) diff --git a/src/gt4py/next/type_system/type_translation.py b/src/gt4py/next/type_system/type_translation.py index 4ff4d3970b..13afc6a313 100644 --- a/src/gt4py/next/type_system/type_translation.py +++ b/src/gt4py/next/type_system/type_translation.py @@ -188,8 +188,14 @@ def canonicalize_type_hint( type_hint = _resolve_type_alias(type_hint) canonical_type = typing.get_origin(type_hint) or type_hint - # In order to distinguish `tuple` from `tuple[()]`, the former returns None here. - args = typing.get_args(type_hint) if typing.get_origin(type_hint) else None + # `None` (rather than `()`) marks an unparametrized annotation, distinguishing bare + # `tuple` / `typing.Tuple` from the empty-tuple annotation `tuple[()]`. Unsubscripted + # `typing` aliases have a truthy origin but no `__args__` attribute. + args = ( + typing.get_args(type_hint) + if typing.get_origin(type_hint) and hasattr(type_hint, "__args__") + else None + ) return canonical_type, args, tuple(extra_args) @@ -222,8 +228,9 @@ def from_type_hint( # Variable-length tuple, e.g. `tuple[int32, ...]` elif isinstance(args, tuple) and len(args) == 2 and args[1] is Ellipsis: return ts.VarArgType(element_type=from_type_hint_same_ns(args[0])) - # Unparametrized annotation, i.e. `tuple` or `tuple[()]` - elif args is None or (isinstance(args, tuple) and len(args) == 0): + # Unparametrized annotation, i.e. bare `tuple` / `typing.Tuple` (the + # empty-tuple annotation `tuple[()]` has `args == ()` and is rejected below) + elif args is None: # TODO(tehrengruber): We use `DeferredType` until we have an actual representation # for a generic type. return ts.DeferredType(constraint=ts.TupleType) diff --git a/tests/next_tests/unit_tests/ffront_tests/test_foast_to_gtir.py b/tests/next_tests/unit_tests/ffront_tests/test_foast_to_gtir.py index 9c4dd08266..e6d7249c36 100644 --- a/tests/next_tests/unit_tests/ffront_tests/test_foast_to_gtir.py +++ b/tests/next_tests/unit_tests/ffront_tests/test_foast_to_gtir.py @@ -531,8 +531,8 @@ def foo(a: gtx.Field[[Edge], float64], b: gtx.Field[[Vertex], float64]): return tuple(2.0 * el for el in (a(V2E), b)) with pytest.raises( - NotImplementedError, - match="fixed-length tuples require all iterable elements to have the same type", + errors.DSLError, + match="fixed-length tuples with differently typed iterable elements are not implemented", ): FieldOperatorParser.apply_to_function(foo) @@ -542,8 +542,8 @@ def foo(a: gtx.Field[[Edge], float64], b: gtx.Field[[Vertex], float64]): return tuple(2.0 * el for el in (a, b)) with pytest.raises( - NotImplementedError, - match="fixed-length tuples require all iterable elements to have the same type", + errors.DSLError, + match="fixed-length tuples with differently typed iterable elements are not implemented", ): FieldOperatorParser.apply_to_function(foo) @@ -553,8 +553,8 @@ def foo(a: gtx.Field[[Edge], float64], b: gtx.Field[[Vertex], float64]): return tuple(2.0 * left + right for left, right in ((a(V2E), b), (b, b))) with pytest.raises( - NotImplementedError, - match="fixed-length tuples require all iterable elements to have the same type", + errors.DSLError, + match="fixed-length tuples with differently typed iterable elements are not implemented", ): FieldOperatorParser.apply_to_function(foo) diff --git a/tests/next_tests/unit_tests/ffront_tests/test_type_deduction.py b/tests/next_tests/unit_tests/ffront_tests/test_type_deduction.py index 22bd1a7a9e..ef1ddb3a3d 100644 --- a/tests/next_tests/unit_tests/ffront_tests/test_type_deduction.py +++ b/tests/next_tests/unit_tests/ffront_tests/test_type_deduction.py @@ -74,6 +74,24 @@ def temp_tuple(a: Field[[TDim], float64], b: Field[[TDim], int64]): ) +def test_tuple_comprehension_unpack_non_tuple(): + def foo(it: tuple[Field[[TDim], float64], ...]): + return tuple(a + b for a, b in it) + + with pytest.raises(errors.DSLError, match=r"Cannot unpack non-iterable"): + _ = FieldOperatorParser.apply_to_function(foo) + + +def test_tuple_comprehension_unpack_too_short_tuple(): + def foo(it: tuple[tuple[Field[[TDim], float64], Field[[TDim], float64]], ...]): + return tuple(a + b + c for a, b, c in it) + + with pytest.raises( + errors.DSLError, match=r"Not enough values to unpack \(expected at least 3, got 2\)" + ): + _ = FieldOperatorParser.apply_to_function(foo) + + def test_adding_bool(): """Expect an error when using arithmetic on bools.""" diff --git a/tests/next_tests/unit_tests/type_system_tests/test_type_info.py b/tests/next_tests/unit_tests/type_system_tests/test_type_info.py index 35c3d2eba1..271688d655 100644 --- a/tests/next_tests/unit_tests/type_system_tests/test_type_info.py +++ b/tests/next_tests/unit_tests/type_system_tests/test_type_info.py @@ -462,3 +462,15 @@ def test_return_type( ) def test_needs_value_extraction(type_spec: ts.TypeSpec, expected: bool): assert type_info.needs_value_extraction(type_spec) is expected + + +def test_is_concretizable_vararg_to_tuple(): + float_type = ts.ScalarType(kind=ts.ScalarKind.FLOAT64) + int_type = ts.ScalarType(kind=ts.ScalarKind.INT32) + vararg = ts.VarArgType(element_type=float_type) + + assert type_info.is_concretizable(vararg, ts.TupleType(types=[float_type, float_type])) + assert not type_info.is_concretizable(vararg, ts.TupleType(types=[int_type])) + # The empty tuple is a valid concretization of any variable-length tuple: there are + # no elements whose type could conflict with the element type. + assert type_info.is_concretizable(vararg, ts.TupleType(types=[])) diff --git a/tests/next_tests/unit_tests/type_system_tests/test_type_translation.py b/tests/next_tests/unit_tests/type_system_tests/test_type_translation.py index fe259f18d0..6e55ca0614 100644 --- a/tests/next_tests/unit_tests/type_system_tests/test_type_translation.py +++ b/tests/next_tests/unit_tests/type_system_tests/test_type_translation.py @@ -220,6 +220,12 @@ def test_invalid_symbol_types(): assert type_translation.from_type_hint(typing.Tuple) == ts.DeferredType(constraint=ts.TupleType) assert type_translation.from_type_hint(tuple) == ts.DeferredType(constraint=ts.TupleType) + # The empty-tuple annotation is not a valid type annotation. + with pytest.raises(ValueError, match="Tuple annotation"): + type_translation.from_type_hint(tuple[()]) + with pytest.raises(ValueError, match="Tuple annotation"): + type_translation.from_type_hint(typing.Tuple[()]) + # Variadic tuples (`tuple[T, ...]`) are now valid — returns a VarArgType. assert type_translation.from_type_hint(tuple[int, ...]) == ts.VarArgType( element_type=ts.ScalarType(kind=ts.ScalarKind.INT64) From 0f11a59b55453e21f0a30e794d5b290b3482e12f Mon Sep 17 00:00:00 2001 From: Till Ehrengruber Date: Tue, 1 Sep 2026 09:10:54 +0000 Subject: [PATCH 06/24] Add empty-tuple comprehension test, xfailed on compiled backends --- .../feature_tests/ffront_tests/test_tuples.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py index 392b900841..39c53e0073 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py @@ -141,6 +141,22 @@ def testee(tracers: tuple[cases.IField, ...], factor: int32) -> tuple[cases.IFie ) +@pytest.mark.uses_tuple_args +def test_var_len_tuple_comprehension_empty(cartesian_case): + # The empty tuple is a valid concretization of a variable-length tuple; the program + # then has nothing to write and must simply run without error. + @gtx.field_operator + def testee(tracers: tuple[cases.IField, ...], factor: int32) -> tuple[cases.IField, ...]: + return tuple(tracer * factor for tracer in tracers) + + if cartesian_case.backend is not None: + pytest.xfail( + "The compiled toolchain cannot execute a program with an empty tuple output: " + "a 'SetAt' for an empty 'out' has no target and no domain to derive." + ) + cases.run(cartesian_case, testee, (), 42, out=()) + + @pytest.mark.uses_tuple_args @pytest.mark.xfail( strict=True, From 681ce77a1a27081c339dde3a66ae6c8e0b2de738 Mon Sep 17 00:00:00 2001 From: Till Ehrengruber Date: Tue, 1 Sep 2026 09:20:55 +0000 Subject: [PATCH 07/24] Add unit test for the generic-tuple constructor type --- .../type_system_tests/test_type_translation.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/tests/next_tests/unit_tests/type_system_tests/test_type_translation.py b/tests/next_tests/unit_tests/type_system_tests/test_type_translation.py index 6e55ca0614..ace88bee3e 100644 --- a/tests/next_tests/unit_tests/type_system_tests/test_type_translation.py +++ b/tests/next_tests/unit_tests/type_system_tests/test_type_translation.py @@ -210,6 +210,23 @@ def test_make_symbol_type_from_typing(value, expected): assert type_translation.from_type_hint(value) == expected +def test_generic_tuple_constructor_type(): + # The type of the `tuple` builtin itself: `type[tuple]` (e.g. the symbol `tuple` + # referenced inside a field operator) maps to a constructor for a generic tuple. + expected = ts.ConstructorType( + definition=ts.FunctionType( + pos_only_args=[ts.DeferredType(constraint=None)], + pos_or_kw_args={}, + kw_only_args={}, + returns=ts.DeferredType(constraint=ts.VarArgType), + ) + ) + assert ( + type_translation.make_constructor_type(ts.DeferredType(constraint=ts.TupleType)) == expected + ) + assert type_translation.from_type_hint(type[tuple]) == expected + + def test_invalid_symbol_types(): # Forward references with pytest.raises(ValueError, match="undefined forward references"): From 5da530ef72fe9306a7e2293b8acd7fd3df67dfef Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Fri, 11 Sep 2026 12:06:49 +0200 Subject: [PATCH 08/24] test[next]: improve coverage for closure vars in tuple comprehension (#2868) --- .../feature_tests/ffront_tests/test_tuples.py | 48 ++++++++++++++++++- .../ffront_tests/test_source_utils.py | 31 ++++++++++-- 2 files changed, 73 insertions(+), 6 deletions(-) diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py index 39c53e0073..5a1b38cab5 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py @@ -10,7 +10,7 @@ import pytest import gt4py.next as gtx -from gt4py.next import broadcast, errors, float64, int32, neighbor_sum, utils as gt_utils +from gt4py.next import broadcast, common, errors, float64, int32, neighbor_sum, utils as gt_utils from next_tests.integration_tests import cases from next_tests.integration_tests.cases import ( @@ -199,6 +199,52 @@ def testee(tracers: tuple[cases.IField, ...], factor: int32) -> tuple[cases.IFie ) +@gtx.field_operator +def _module_level_inner(tracer: cases.IField, factor: int32) -> cases.IField: + return tracer * factor + + +@pytest.mark.uses_tuple_args +def test_tuple_comprehension_module_level_fo(cartesian_case): + @gtx.field_operator + def testee(tracers: tuple[cases.IField, ...], factor: int32) -> tuple[cases.IField, ...]: + return tuple(_module_level_inner(tracer, factor) for tracer in tracers) + + cases.verify_with_default_data( + cartesian_case, testee, ref=lambda t, f: tuple(el * f for el in t) + ) + + +@pytest.mark.uses_tuple_args +@pytest.mark.uses_unstructured_shift +def test_tuple_comprehension_module_level_builtin(unstructured_case): + @gtx.field_operator + def testee(tracers: tuple[cases.EField, ...]) -> tuple[cases.VField, ...]: + return tuple(neighbor_sum(tracer(V2E), axis=V2EDim) for tracer in tracers) + + v2e = unstructured_case.offset_provider["V2E"].asnumpy() + valid = v2e != common._DEFAULT_SKIP_VALUE + cases.verify_with_default_data( + unstructured_case, + testee, + ref=lambda t: tuple(np.sum(el[v2e], axis=1, where=valid) for el in t), + ) + + +@pytest.mark.uses_tuple_args +def test_tuple_comprehension_module_level_fo_used_outside_too(cartesian_case): + @gtx.field_operator + def testee( + tracers: tuple[cases.IField, ...], a: cases.IField, factor: int32 + ) -> tuple[cases.IField, ...]: + unrelated = _module_level_inner(a, factor) + return tuple(_module_level_inner(tracer, factor) for tracer in tracers) + + cases.verify_with_default_data( + cartesian_case, testee, ref=lambda t, a, f: tuple(el * f for el in t) + ) + + @pytest.mark.uses_tuple_args def test_nested_tuple_comprehension(cartesian_case): @gtx.field_operator diff --git a/tests/next_tests/unit_tests/ffront_tests/test_source_utils.py b/tests/next_tests/unit_tests/ffront_tests/test_source_utils.py index 5735b4c25b..9a08544d91 100644 --- a/tests/next_tests/unit_tests/ffront_tests/test_source_utils.py +++ b/tests/next_tests/unit_tests/ffront_tests/test_source_utils.py @@ -10,17 +10,19 @@ A generator expression, lambda or nested function compiles to its own code object, so names referenced only inside it appear in that object's `co_names` and never in -the enclosing function's. Collecting closure variables from the enclosing function -alone therefore misses them; `get_closure_vars_from_function` takes the global names -from the compiler's symbol table of the source instead, which covers nested scopes. +the enclosing function's. `get_closure_vars_from_function` therefore takes the global +names from the compiler's symbol table of the source, which covers nested scopes. Free variables need no such treatment: the enclosing code object carries a cell for them. """ +import pytest + import gt4py.next as gtx from gt4py.next import Dims, Dimension, float64, neighbor_sum from gt4py.next.ffront import source_utils +from gt4py.next.ffront.func_to_foast import FieldOperatorParser from gt4py.next.ffront.source_utils import get_closure_vars_from_function @@ -46,6 +48,27 @@ def _module_level_operator(tracers: tuple[CField, ...], factor: float64) -> tupl return tuple(scale(t, factor) for t in tracers) +@pytest.mark.parametrize( + "func", [_builtin_and_offset, _module_level_operator], ids=["builtin_and_offset", "operator"] +) +def test_module_level_name_used_only_in_comprehension(func): + """Module-level names are reachable from inside a comprehension body.""" + parsed = FieldOperatorParser.apply_to_function(func) + assert parsed.type is not None + + +def test_free_variables_still_resolve_through_the_parser(): + @gtx.field_operator + def local_scale(f: CField, factor: float64) -> CField: + return f * factor + + def uses_freevar(tracers: tuple[CField, ...], factor: float64) -> tuple[CField, ...]: + return tuple(local_scale(t, factor) for t in tracers) + + parsed = FieldOperatorParser.apply_to_function(uses_freevar) + assert parsed.type is not None + + def test_names_are_collected_from_the_nested_code_object(): """The names live in the generator expression's code object, not the function's.""" assert "neighbor_sum" not in _builtin_and_offset.__code__.co_names @@ -62,8 +85,6 @@ def test_comprehension_target_is_not_collected(): def test_free_variables_still_resolve(): - """The path that already worked, kept as a guard.""" - @gtx.field_operator def local_scale(f: CField, factor: float64) -> CField: return f * factor From c1ccbad552e2586c695b8425c87e2f3e67e43cfe Mon Sep 17 00:00:00 2001 From: Till Ehrengruber Date: Tue, 1 Sep 2026 09:43:32 +0000 Subject: [PATCH 09/24] Treat operators with variable-length tuple params as generic in '_is_generic' --- src/gt4py/next/otf/compiled_program.py | 13 ++++++++++-- .../feature_tests/ffront_tests/test_tuples.py | 21 +++++++++++++++++++ 2 files changed, 32 insertions(+), 2 deletions(-) diff --git a/src/gt4py/next/otf/compiled_program.py b/src/gt4py/next/otf/compiled_program.py index 0b49f9ec0c..4d69e11d8c 100644 --- a/src/gt4py/next/otf/compiled_program.py +++ b/src/gt4py/next/otf/compiled_program.py @@ -507,13 +507,22 @@ def _is_generic(self) -> bool: Is the operator or program generic in the sense that it can be called for different argument types. - Right now this is only the case for scan operators. + Right now this is the case for scan operators and for parameters containing a + variable-length tuple, whose concrete length is only known at the call. """ + # TODO(tehrengruber): This concept does not exist elsewhere and is not properly reflected # in the type system. For now we just use `DeferredType` to communicate between # here and `type_info.type_in_program_context`. + def is_generic_type(t: ts.TypeSpec) -> bool: + if isinstance(t, (ts.DeferredType, ts.VarArgType)): + return True + if isinstance(t, ts.TupleType): + return any(is_generic_type(el) for el in t.types) + return False + return any( - isinstance(t, ts.DeferredType) + is_generic_type(t) for t in itertools.chain( self.program_type.definition.pos_only_args, self.program_type.definition.pos_or_kw_args.values(), diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py index 5a1b38cab5..fd3aab9a3d 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py @@ -157,6 +157,27 @@ def testee(tracers: tuple[cases.IField, ...], factor: int32) -> tuple[cases.IFie cases.run(cartesian_case, testee, (), 42, out=()) +@pytest.mark.uses_tuple_args +def test_var_len_tuple_comprehension_different_lengths(cartesian_case): + # An operator with variable-length tuple parameters is generic: every call must be + # specialized (and compiled) for the concrete tuple length of its arguments. + @gtx.field_operator + def testee(tracers: tuple[cases.IField, ...], factor: int32) -> tuple[cases.IField, ...]: + return tuple(tracer * factor for tracer in tracers) + + for length in (3, 2): + tracers = cases.allocate(cartesian_case, testee, "tracers")()[:length] + out = cases.allocate(cartesian_case, testee, cases.RETURN)()[:length] + cases.verify( + cartesian_case, + testee, + tracers, + 42, + out=out, + ref=tuple(t.asnumpy() * 42 for t in tracers), + ) + + @pytest.mark.uses_tuple_args @pytest.mark.xfail( strict=True, From 68de5f4b7e81ee13b8b9d1af2dea028d3c6c9de5 Mon Sep 17 00:00:00 2001 From: Till Ehrengruber Date: Mon, 21 Sep 2026 11:53:23 +0000 Subject: [PATCH 10/24] Address review comments (unpack arity, nested vararg concretization, target validation, annotation errors, pretty printer) --- .../ffront/foast_passes/type_deduction.py | 19 +++++++++++ src/gt4py/next/ffront/foast_pretty_printer.py | 8 +++-- src/gt4py/next/ffront/func_to_foast.py | 15 ++++++++- src/gt4py/next/type_system/type_info.py | 13 +++++--- .../next/type_system/type_translation.py | 12 +++++-- .../ffront_tests/test_foast_pretty_printer.py | 12 +++++++ .../feature_tests/ffront_tests/test_tuples.py | 23 +++++++++++++ .../ffront_tests/test_func_to_foast.py | 32 +++++++++++++++++++ .../ffront_tests/test_type_deduction.py | 12 +++++++ .../type_system_tests/test_type_info.py | 14 ++++++++ .../test_type_translation.py | 6 ++++ 11 files changed, 157 insertions(+), 9 deletions(-) diff --git a/src/gt4py/next/ffront/foast_passes/type_deduction.py b/src/gt4py/next/ffront/foast_passes/type_deduction.py index c616a7ea10..d2c941c2d7 100644 --- a/src/gt4py/next/ffront/foast_passes/type_deduction.py +++ b/src/gt4py/next/ffront/foast_passes/type_deduction.py @@ -758,6 +758,25 @@ def _deduce_tuple_comprehension_target_type( replaces that by the constituent of `element_type` at the symbol's unpacking position. """ + def first_symbol(target_el: MaybeNestedInTuple[foast.Symbol]) -> foast.Symbol: + return first_symbol(target_el[0]) if isinstance(target_el, tuple) else target_el + + def ensure_all_values_consumed( + target_el: MaybeNestedInTuple[foast.Symbol], type_: ts.TypeSpec + ) -> None: + # An unpacking pattern with fewer names than the element has values would + # silently drop the excess values; reject it like the equivalent assignment. + if isinstance(target_el, tuple) and isinstance(type_, ts.TupleType): + if len(type_.types) > len(target_el): + raise errors.DSLError( + first_symbol(target_el).location, + f"Too many values to unpack (expected {len(target_el)}).", + ) + for el, el_type in zip(target_el, type_.types): + ensure_all_values_consumed(el, el_type) + + ensure_all_values_consumed(target, element_type) + @tree_map(with_path_arg=True) def process_target(target_el: foast.Symbol, path: tuple[int, ...]) -> foast.Symbol: type_ = element_type diff --git a/src/gt4py/next/ffront/foast_pretty_printer.py b/src/gt4py/next/ffront/foast_pretty_printer.py index 79826983a3..8c4203be6d 100644 --- a/src/gt4py/next/ffront/foast_pretty_printer.py +++ b/src/gt4py/next/ffront/foast_pretty_printer.py @@ -123,9 +123,13 @@ def apply(cls, node: foast.LocatedNode, **kwargs: Any) -> str: # type: ignore[o def visit_TupleComprehensionMapper( self, node: foast.TupleComprehensionMapper, **kwargs: Any ) -> str: + def format_target(target: Any) -> str: + if isinstance(target, tuple): + return f"({', '.join(format_target(el) for el in target)})" + return self.visit(target, **kwargs) + element_expr = self.visit(node.element_expr, **kwargs) - target = self.visit(node.target, **kwargs) - return f"{element_expr} for {target}" + return f"{element_expr} for {format_target(node.target)}" def visit_TupleComprehension(self, node: foast.TupleComprehension, **kwargs: Any) -> str: mapper = self.visit(node.inner, **kwargs) diff --git a/src/gt4py/next/ffront/func_to_foast.py b/src/gt4py/next/ffront/func_to_foast.py index b210027d5d..da30ab9596 100644 --- a/src/gt4py/next/ffront/func_to_foast.py +++ b/src/gt4py/next/ffront/func_to_foast.py @@ -594,10 +594,23 @@ def visit_Call(self, node: ast.Call, **kwargs: Any) -> foast.Call | foast.TupleC ], ) + seen_target_names: set[str] = set() + def parse_target(target: ast.expr) -> MaybeNestedInTuple[foast.DataSymbol]: if isinstance(target, ast.Tuple): return tuple(parse_target(el) for el in target.elts) - assert isinstance(target, ast.Name) and isinstance(target.ctx, ast.Store) + if not (isinstance(target, ast.Name) and isinstance(target.ctx, ast.Store)): + raise errors.UnsupportedPythonFeatureError( + self.get_location(target), + "generator expression targets other than names " + "and (nested) tuples of names", + ) + if target.id in seen_target_names: + raise errors.DSLError( + self.get_location(target), + f"Duplicate name '{target.id}' in generator expression target.", + ) + seen_target_names.add(target.id) return foast.DataSymbol( id=target.id, location=self.get_location(target), diff --git a/src/gt4py/next/type_system/type_info.py b/src/gt4py/next/type_system/type_info.py index 9eb37301c8..55997febcf 100644 --- a/src/gt4py/next/type_system/type_info.py +++ b/src/gt4py/next/type_system/type_info.py @@ -555,12 +555,17 @@ def is_concretizable(symbol_type: ts.TypeSpec, to_type: ts.TypeSpec) -> bool: if isinstance(symbol_type, ts.VarArgType) and isinstance(to_type, ts.TupleType): # The empty tuple is a valid concretization of any variable-length tuple: there # are no elements whose type could conflict with the element type. - if len(to_type.types) == 0 or ( + return len(to_type.types) == 0 or ( all(type_ == to_type.types[0] for type_ in to_type.types) and is_concretizable(symbol_type.element_type, to_type.types[0]) - ): - return True - elif is_concrete(symbol_type): + ) + if isinstance(symbol_type, ts.TupleType) and isinstance(to_type, ts.TupleType): + # Recurse elementwise so that e.g. a variable-length tuple nested inside a + # fixed-length tuple can be concretized. + return len(symbol_type.types) == len(to_type.types) and all( + is_concretizable(el, to_el) for el, to_el in zip(symbol_type.types, to_type.types) + ) + if is_concrete(symbol_type): return symbol_type == to_type return False diff --git a/src/gt4py/next/type_system/type_translation.py b/src/gt4py/next/type_system/type_translation.py index 13afc6a313..1070ccdbd5 100644 --- a/src/gt4py/next/type_system/type_translation.py +++ b/src/gt4py/next/type_system/type_translation.py @@ -223,11 +223,19 @@ def from_type_hint( and not any(arg is Ellipsis for arg in args) ): tuple_types = [from_type_hint_same_ns(arg) for arg in args] - assert all(isinstance(elem, ts.DataType) for elem in tuple_types) + if not all(isinstance(elem, ts.DataType) for elem in tuple_types): + raise ValueError( + f"Tuple annotation '{type_hint}' contains invalid element types." + ) return ts.TupleType(types=tuple_types) # Variable-length tuple, e.g. `tuple[int32, ...]` elif isinstance(args, tuple) and len(args) == 2 and args[1] is Ellipsis: - return ts.VarArgType(element_type=from_type_hint_same_ns(args[0])) + element_type = from_type_hint_same_ns(args[0]) + if not isinstance(element_type, ts.DataType): + raise ValueError( + f"Tuple annotation '{type_hint}' contains an invalid element type." + ) + return ts.VarArgType(element_type=element_type) # Unparametrized annotation, i.e. bare `tuple` / `typing.Tuple` (the # empty-tuple annotation `tuple[()]` has `args == ()` and is rejected below) elif args is None: diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_foast_pretty_printer.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_foast_pretty_printer.py index 8ba6f3091a..734d734a91 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_foast_pretty_printer.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_foast_pretty_printer.py @@ -84,3 +84,15 @@ def scan(inp: int32) -> int32: ).strip() assert pretty_format(scan.foast_stage.foast_node) == expected + + +def test_tuple_comprehension_multi_target(): + I = Dimension("I") + + @field_operator + def foo( + b: tuple[tuple[Field[[I], int64], Field[[I], int64]], ...], + ) -> tuple[Field[[I], int64], ...]: + return tuple(x + y for x, y in b) + + assert "for (x, y) in b" in pretty_format(foo.foast_stage.foast_node) diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py index fd3aab9a3d..354451e45f 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py @@ -157,6 +157,29 @@ def testee(tracers: tuple[cases.IField, ...], factor: int32) -> tuple[cases.IFie cases.run(cartesian_case, testee, (), 42, out=()) +@pytest.mark.uses_tuple_args +@pytest.mark.uses_tuple_returns +def test_var_len_tuple_comprehension_in_fixed_tuple_return(cartesian_case): + # A comprehension result nested inside a fixed-length return tuple, called directly. + @gtx.field_operator + def testee( + tracers: tuple[cases.IField, ...], rho: cases.IField + ) -> tuple[tuple[cases.IField, ...], cases.IField]: + return tuple(tracer * 2 for tracer in tracers), rho + + tracers = cases.allocate(cartesian_case, testee, "tracers")() + rho = cases.allocate(cartesian_case, testee, "rho")() + out = cases.allocate(cartesian_case, testee, cases.RETURN)() + cases.verify( + cartesian_case, + testee, + tracers, + rho, + out=out, + ref=(tuple(t.asnumpy() * 2 for t in tracers), rho.asnumpy()), + ) + + @pytest.mark.uses_tuple_args def test_var_len_tuple_comprehension_different_lengths(cartesian_case): # An operator with variable-length tuple parameters is generic: every call must be diff --git a/tests/next_tests/unit_tests/ffront_tests/test_func_to_foast.py b/tests/next_tests/unit_tests/ffront_tests/test_func_to_foast.py index 39c7a1cb8b..7847444535 100644 --- a/tests/next_tests/unit_tests/ffront_tests/test_func_to_foast.py +++ b/tests/next_tests/unit_tests/ffront_tests/test_func_to_foast.py @@ -509,6 +509,38 @@ def tuple_index_failure( _ = FieldOperatorParser.apply_to_function(tuple_index_failure) +def test_tuple_compr_list_target_failure(): + def testee(arg: tuple[tuple[int32, int32], ...]): + return tuple(p + q for [p, q] in arg) + + with pytest.raises(errors.UnsupportedPythonFeatureError, match=r"generator expression targets"): + _ = FieldOperatorParser.apply_to_function(testee) + + +def test_tuple_compr_starred_target_failure(): + def testee(arg: tuple[tuple[int32, int32], ...]): + return tuple(p for p, *q in arg) + + with pytest.raises(errors.UnsupportedPythonFeatureError, match=r"generator expression targets"): + _ = FieldOperatorParser.apply_to_function(testee) + + +def test_tuple_compr_subscript_target_failure(): + def testee(arg: tuple[int32, ...], x: tuple[int32, int32]): + return tuple(x[0] for x[0] in arg) + + with pytest.raises(errors.UnsupportedPythonFeatureError, match=r"generator expression targets"): + _ = FieldOperatorParser.apply_to_function(testee) + + +def test_tuple_compr_duplicate_target_name_failure(): + def testee(arg: tuple[tuple[int32, int32], ...]): + return tuple(p for p, p in arg) + + with pytest.raises(errors.DSLError, match=r"Duplicate name 'p'"): + _ = FieldOperatorParser.apply_to_function(testee) + + def test_tuple_compr_non_tuple_iterable_failure(): def testee(arg: float): return tuple(_ for _ in arg) diff --git a/tests/next_tests/unit_tests/ffront_tests/test_type_deduction.py b/tests/next_tests/unit_tests/ffront_tests/test_type_deduction.py index 1585b433b8..4eae9abf48 100644 --- a/tests/next_tests/unit_tests/ffront_tests/test_type_deduction.py +++ b/tests/next_tests/unit_tests/ffront_tests/test_type_deduction.py @@ -82,6 +82,18 @@ def foo(it: tuple[Field[[TDim], float64], ...]): _ = FieldOperatorParser.apply_to_function(foo) +def test_tuple_comprehension_unpack_too_many_values(): + def foo( + it: tuple[ + tuple[Field[[TDim], float64], Field[[TDim], float64], Field[[TDim], float64]], ... + ], + ): + return tuple(a + b for a, b in it) + + with pytest.raises(errors.DSLError, match=r"Too many values to unpack \(expected 2\)"): + _ = FieldOperatorParser.apply_to_function(foo) + + def test_tuple_comprehension_unpack_too_short_tuple(): def foo(it: tuple[tuple[Field[[TDim], float64], Field[[TDim], float64]], ...]): return tuple(a + b + c for a, b, c in it) diff --git a/tests/next_tests/unit_tests/type_system_tests/test_type_info.py b/tests/next_tests/unit_tests/type_system_tests/test_type_info.py index 3c806f4242..019d4cf648 100644 --- a/tests/next_tests/unit_tests/type_system_tests/test_type_info.py +++ b/tests/next_tests/unit_tests/type_system_tests/test_type_info.py @@ -491,3 +491,17 @@ def test_promote_lists(): type_info.promote(v2e_list, ts.ListType(element_type=int32, offset_type=V2EDim)) with pytest.raises(ValueError, match="non-lists"): type_info.promote(v2e_list, float64) + + +def test_is_concretizable_tuple_with_nested_vararg(): + float_type = ts.ScalarType(kind=ts.ScalarKind.FLOAT64) + field_type = ts.FieldType(dims=[], dtype=float_type) + nested = ts.TupleType(types=[ts.VarArgType(element_type=field_type), field_type]) + + assert type_info.is_concretizable( + nested, ts.TupleType(types=[ts.TupleType(types=[field_type, field_type]), field_type]) + ) + assert not type_info.is_concretizable( + nested, ts.TupleType(types=[ts.TupleType(types=[float_type]), field_type]) + ) + assert not type_info.is_concretizable(nested, ts.TupleType(types=[field_type])) diff --git a/tests/next_tests/unit_tests/type_system_tests/test_type_translation.py b/tests/next_tests/unit_tests/type_system_tests/test_type_translation.py index ace88bee3e..20a772a1c9 100644 --- a/tests/next_tests/unit_tests/type_system_tests/test_type_translation.py +++ b/tests/next_tests/unit_tests/type_system_tests/test_type_translation.py @@ -243,6 +243,12 @@ def test_invalid_symbol_types(): with pytest.raises(ValueError, match="Tuple annotation"): type_translation.from_type_hint(typing.Tuple[()]) + # Invalid element annotations are reported as ValueError, not internal errors. + with pytest.raises(ValueError, match="invalid element type"): + type_translation.from_type_hint(tuple[tuple, int]) + with pytest.raises(ValueError, match="invalid element type"): + type_translation.from_type_hint(tuple[tuple, ...]) + # Variadic tuples (`tuple[T, ...]`) are now valid — returns a VarArgType. assert type_translation.from_type_hint(tuple[int, ...]) == ts.VarArgType( element_type=ts.ScalarType(kind=ts.ScalarKind.INT64) From 840b85ea95b09b96f17367abe5f3b1db09322902 Mon Sep 17 00:00:00 2001 From: Till Ehrengruber Date: Mon, 21 Sep 2026 15:46:38 +0000 Subject: [PATCH 11/24] Fold the unpack-arity check into 'process_target' --- .../ffront/foast_passes/type_deduction.py | 33 +++++++------------ .../ffront_tests/test_type_deduction.py | 2 +- 2 files changed, 13 insertions(+), 22 deletions(-) diff --git a/src/gt4py/next/ffront/foast_passes/type_deduction.py b/src/gt4py/next/ffront/foast_passes/type_deduction.py index d2c941c2d7..899a18fcd7 100644 --- a/src/gt4py/next/ffront/foast_passes/type_deduction.py +++ b/src/gt4py/next/ffront/foast_passes/type_deduction.py @@ -758,40 +758,31 @@ def _deduce_tuple_comprehension_target_type( replaces that by the constituent of `element_type` at the symbol's unpacking position. """ - def first_symbol(target_el: MaybeNestedInTuple[foast.Symbol]) -> foast.Symbol: - return first_symbol(target_el[0]) if isinstance(target_el, tuple) else target_el - - def ensure_all_values_consumed( - target_el: MaybeNestedInTuple[foast.Symbol], type_: ts.TypeSpec - ) -> None: - # An unpacking pattern with fewer names than the element has values would - # silently drop the excess values; reject it like the equivalent assignment. - if isinstance(target_el, tuple) and isinstance(type_, ts.TupleType): - if len(type_.types) > len(target_el): - raise errors.DSLError( - first_symbol(target_el).location, - f"Too many values to unpack (expected {len(target_el)}).", - ) - for el, el_type in zip(target_el, type_.types): - ensure_all_values_consumed(el, el_type) - - ensure_all_values_consumed(target, element_type) - @tree_map(with_path_arg=True) def process_target(target_el: foast.Symbol, path: tuple[int, ...]) -> foast.Symbol: + # Walk the unpacking pattern and the element type side-by-side so that the + # number of names and values can be compared at every nesting level. type_ = element_type + target_tree: MaybeNestedInTuple[foast.Symbol] = target for i in path: + assert isinstance(target_tree, tuple) if not isinstance(type_, ts.TupleType): raise errors.DSLError( target_el.location, f"Cannot unpack non-iterable '{type_}' object." ) - if len(type_.types) <= i: + if len(type_.types) > len(target_tree): + raise errors.DSLError( + target_el.location, + f"Too many values to unpack (expected {len(target_tree)}).", + ) + if len(type_.types) < len(target_tree): raise errors.DSLError( target_el.location, - f"Not enough values to unpack (expected at least {i + 1}, " + f"Not enough values to unpack (expected {len(target_tree)}, " f"got {len(type_.types)}).", ) type_ = type_.types[i] + target_tree = target_tree[i] return self.visit(target_el, refine_type=type_, **kwargs) return process_target(target) diff --git a/tests/next_tests/unit_tests/ffront_tests/test_type_deduction.py b/tests/next_tests/unit_tests/ffront_tests/test_type_deduction.py index 4eae9abf48..75727f103a 100644 --- a/tests/next_tests/unit_tests/ffront_tests/test_type_deduction.py +++ b/tests/next_tests/unit_tests/ffront_tests/test_type_deduction.py @@ -99,7 +99,7 @@ def foo(it: tuple[tuple[Field[[TDim], float64], Field[[TDim], float64]], ...]): return tuple(a + b + c for a, b, c in it) with pytest.raises( - errors.DSLError, match=r"Not enough values to unpack \(expected at least 3, got 2\)" + errors.DSLError, match=r"Not enough values to unpack \(expected 3, got 2\)" ): _ = FieldOperatorParser.apply_to_function(foo) From bbb50991d30f565143a8b32de40094ebb7a32e4d Mon Sep 17 00:00:00 2001 From: Till Ehrengruber Date: Tue, 22 Sep 2026 19:22:10 +0200 Subject: [PATCH 12/24] Apply pre-commit ruff formatting --- .../next_tests/unit_tests/ffront_tests/test_type_deduction.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/next_tests/unit_tests/ffront_tests/test_type_deduction.py b/tests/next_tests/unit_tests/ffront_tests/test_type_deduction.py index 75727f103a..1ddd7d6d1d 100644 --- a/tests/next_tests/unit_tests/ffront_tests/test_type_deduction.py +++ b/tests/next_tests/unit_tests/ffront_tests/test_type_deduction.py @@ -98,9 +98,7 @@ def test_tuple_comprehension_unpack_too_short_tuple(): def foo(it: tuple[tuple[Field[[TDim], float64], Field[[TDim], float64]], ...]): return tuple(a + b + c for a, b, c in it) - with pytest.raises( - errors.DSLError, match=r"Not enough values to unpack \(expected 3, got 2\)" - ): + with pytest.raises(errors.DSLError, match=r"Not enough values to unpack \(expected 3, got 2\)"): _ = FieldOperatorParser.apply_to_function(foo) From 97576ae4f4afd4cb23e276dfe642cc36c3674bbe Mon Sep 17 00:00:00 2001 From: Till Ehrengruber Date: Wed, 7 Oct 2026 09:05:45 +0200 Subject: [PATCH 13/24] Add TODO about renaming 'VarArgType' --- src/gt4py/next/type_system/type_specifications.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/gt4py/next/type_system/type_specifications.py b/src/gt4py/next/type_system/type_specifications.py index 1016eba34e..c708dc6c77 100644 --- a/src/gt4py/next/type_system/type_specifications.py +++ b/src/gt4py/next/type_system/type_specifications.py @@ -148,6 +148,7 @@ def __len__(self) -> int: return len(self.types) +# TODO(tehrengruber): Rename to var length tuple or similar class VarArgType(DataType): """Represents a variable number of arguments of the same type.""" From 343ec9b3224d405734ea81c7df378de829bf551b Mon Sep 17 00:00:00 2001 From: Till Ehrengruber Date: Wed, 7 Oct 2026 09:40:27 +0200 Subject: [PATCH 14/24] Fix imports after 'eve.extended_typing' was renamed to 'eve.xtyping' (silent merge conflict) --- src/gt4py/next/ffront/field_operator_ast.py | 2 +- src/gt4py/next/ffront/foast_passes/type_deduction.py | 2 +- src/gt4py/next/ffront/func_to_foast.py | 3 ++- 3 files changed, 4 insertions(+), 3 deletions(-) diff --git a/src/gt4py/next/ffront/field_operator_ast.py b/src/gt4py/next/ffront/field_operator_ast.py index f95f9f285a..ba99476d23 100644 --- a/src/gt4py/next/ffront/field_operator_ast.py +++ b/src/gt4py/next/ffront/field_operator_ast.py @@ -12,9 +12,9 @@ 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.eve.xtyping import MaybeNestedInTuple from gt4py.next import utils from gt4py.next.ffront import dialect_ast_enums, type_specifications as ts_ffront from gt4py.next.type_system import type_specifications as ts diff --git a/src/gt4py/next/ffront/foast_passes/type_deduction.py b/src/gt4py/next/ffront/foast_passes/type_deduction.py index a4a6b6711b..bcea684b78 100644 --- a/src/gt4py/next/ffront/foast_passes/type_deduction.py +++ b/src/gt4py/next/ffront/foast_passes/type_deduction.py @@ -12,7 +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.eve.xtyping import MaybeNestedInTuple from gt4py.next import common, errors from gt4py.next.common import Dimension, DimensionKind, promote_dims from gt4py.next.ffront import ( diff --git a/src/gt4py/next/ffront/func_to_foast.py b/src/gt4py/next/ffront/func_to_foast.py index da30ab9596..1820c3203b 100644 --- a/src/gt4py/next/ffront/func_to_foast.py +++ b/src/gt4py/next/ffront/func_to_foast.py @@ -11,9 +11,10 @@ import ast import textwrap import typing +from typing import Any import gt4py.eve as eve -from gt4py.eve.extended_typing import Any, MaybeNestedInTuple +from gt4py.eve.xtyping import MaybeNestedInTuple from gt4py.next import errors from gt4py.next.ffront import ( dialect_ast_enums, From effbb2dd01cfaf9c5597968d25c34aaab2ab986f Mon Sep 17 00:00:00 2001 From: Till Ehrengruber Date: Wed, 14 Oct 2026 23:54:46 +0200 Subject: [PATCH 15/24] feat[next]: variable-length tuple ('tuple[T, ...]') support for programs Hand-written programs can now declare parameters as variable-length tuples; the concrete length is taken from the arguments of the call: - Specialize PAST program parameter types to the concrete argument types in 'past_to_gtir' ('_specialize_program_params'). Generated wrapper programs for direct operator calls now carry the operator's declared (possibly generic) types and are specialized the same way, before their type deduction, as e.g. scan operator parameter types are only known at the call. - Lower the 'out' argument length-agnostically: the target is the plain lowered 'out' expression and the domain tree is built with the new 'tree_map_type_into_expr' helper, using the 'map_tuple' builtin at variable-length nodes. - Lower the 'domain' argument independently of 'out'; a single domain may now apply to an entire tuple subtree of the target (documented in 'itir.SetAt', promoted in the dace backend via 'equalize_tuple_structure'). This revives the shallow domain structure that was the norm before #2225. - Make 'tree_map_type' and 'primitive_constituents' traverse 'VarArgType' (single representative element, addressed by the 'VAR_ARG_ELEMENT' path key). - Include the program parameter types in 'itir.Program.__hash__' as length variants of a program only differ in their parameter types. - Reject subscripts on tuple arguments, per-element domain tuples for variable-length outputs, and AOT compilation of generic programs with proper error messages; validate domain dimensions against the out field constituents at decoration time. - Test infrastructure: 'varargs_length' argument for 'cases.allocate' instead of slicing default-length allocations. --- src/gt4py/next/ffront/decorator.py | 2 + src/gt4py/next/ffront/foast_to_past.py | 46 ++-- .../next/ffront/past_passes/type_deduction.py | 32 ++- src/gt4py/next/ffront/past_to_itir.py | 205 +++++++++--------- src/gt4py/next/ffront/type_info.py | 4 +- src/gt4py/next/iterator/ir.py | 18 +- src/gt4py/next/iterator/ir_utils/misc.py | 88 +++++++- src/gt4py/next/otf/compiled_program.py | 6 + .../runners/dace/lowering/gtir_to_sdfg.py | 4 + src/gt4py/next/type_system/type_info.py | 58 ++++- .../next/type_system/type_specifications.py | 2 + src/gt4py/next/utils.py | 12 +- tests/next_tests/integration_tests/cases.py | 43 +++- .../ffront_tests/test_program.py | 14 +- .../feature_tests/ffront_tests/test_tuples.py | 152 ++++++++++++- .../ffront_tests/test_func_to_past.py | 76 +++++++ 16 files changed, 600 insertions(+), 162 deletions(-) diff --git a/src/gt4py/next/ffront/decorator.py b/src/gt4py/next/ffront/decorator.py index bf63aa692a..09b890413b 100644 --- a/src/gt4py/next/ffront/decorator.py +++ b/src/gt4py/next/ffront/decorator.py @@ -214,6 +214,8 @@ def _field_domain_descriptor_mapping_from_func_type(func_type: ts.FunctionType) for name, type_ in param_types.items(): for el_type_, path in type_info.primitive_constituents(type_, with_path_arg=True): if isinstance(el_type_, ts.FieldType): + if any(isinstance(idx, type_info.VarArgElementKey) for idx in path): + raise NotImplementedError() path_as_expr = "".join(f"[{idx}]" for idx in path) static_domain_args.append(f"{name}{path_as_expr}") return static_domain_args diff --git a/src/gt4py/next/ffront/foast_to_past.py b/src/gt4py/next/ffront/foast_to_past.py index 1bc91ad01b..fd51787e9b 100644 --- a/src/gt4py/next/ffront/foast_to_past.py +++ b/src/gt4py/next/ffront/foast_to_past.py @@ -12,6 +12,7 @@ from gt4py.next.ffront import ( dialect_ast_enums, foast_to_gtir, + past_to_itir, program_ast as past, stages as ffront_stages, type_info as ffront_type_info, @@ -21,7 +22,7 @@ from gt4py.next.ffront.stages import ConcreteFOASTOperatorDef, ConcretePASTProgramDef from gt4py.next.iterator import ir as itir from gt4py.next.otf import workflow -from gt4py.next.type_system import type_info, type_specifications as ts +from gt4py.next.type_system import type_specifications as ts @dataclasses.dataclass(frozen=True) @@ -99,40 +100,28 @@ def __call__(self, inp: ConcreteFOASTOperatorDef) -> ConcretePASTProgramDef: # of arg and kwarg types # TODO(tehrengruber): check foast operator has no out argument that clashes # with the out argument of the program we generate here. - arg_types, kwarg_types = inp.args.args, inp.args.kwargs - assert not kwarg_types - - type_ = inp.data.foast_node.type loc = inp.data.foast_node.location assert isinstance(inp.data.foast_node.type, ts.CallableType) partial_program_type = ffront_type_info.type_in_program_context(inp.data.foast_node.type) assert isinstance(partial_program_type, ts_ffront.ProgramType) - args_names = [ - *partial_program_type.definition.pos_only_args, - *partial_program_type.definition.pos_or_kw_args.keys(), - *partial_program_type.definition.kw_only_args.keys(), - ] - assert isinstance(type_, ts.CallableType) - return_type = type_info.return_type( - type_, with_args=list(arg_types), with_kwargs=kwarg_types + param_types: tuple[tuple[str, ts.TypeSpec], ...] = ( + *( + (f"__pos_arg{i}", type_) + for i, type_ in enumerate(partial_program_type.definition.pos_only_args) + ), + *partial_program_type.definition.pos_or_kw_args.items(), + *partial_program_type.definition.kw_only_args.items(), ) - # 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" + assert param_types[-1][0] == "out" params_decl: list[past.Symbol] = [ past.DataSymbol( id=name, - type=arg_type, + type=type_, # note: this might be a generic type like VarArgType namespace=dialect_ast_enums.Namespace.LOCAL, location=loc, ) - for name, arg_type in zip( - args_names, - arg_types, - strict=True, - ) + for name, type_ in param_types ] params_ref = [past.Name(id=pdecl.id, location=loc) for pdecl in params_decl[:-1]] out_ref = past.Name(id="out", location=loc) @@ -166,6 +155,17 @@ def __call__(self, inp: ConcreteFOASTOperatorDef) -> ConcretePASTProgramDef: closure_vars=closure_symbols, location=loc, ) + # Specialize the parameter types to the concrete argument types of the call before + # the type deduction: the declared types may be generic — e.g. `DeferredType` for + # scan operator parameters, which are only promoted to fields at the call — and the + # body call can not be type checked against them. The specialization happens again + # later, in `past_to_gtir`, which also covers hand-written programs that are not + # specialized here. With support for generic types in the program type deduction + # this step would become obsolete, as the generic variant could be type checked + # directly. + untyped_past_node = past_to_itir._specialize_program_params( + untyped_past_node, inp.args.args, inp.args.kwargs + ) untyped_past_node = closure_var_type_deduction.ClosureVarTypeDeduction.apply( untyped_past_node, fieldop_itir_closure_vars, # type: ignore[arg-type] diff --git a/src/gt4py/next/ffront/past_passes/type_deduction.py b/src/gt4py/next/ffront/past_passes/type_deduction.py index fcc75b7374..d282ccf7d7 100644 --- a/src/gt4py/next/ffront/past_passes/type_deduction.py +++ b/src/gt4py/next/ffront/past_passes/type_deduction.py @@ -57,6 +57,21 @@ def _validate_domain_out( f"Only 'Dimension' allowed in domain dictionary keys, got '{dim}' which is of type '{dim.type}'." ) + dict_dims = [cast(ts.DimensionType, dim.type).dim for dim in dom.keys_] + + # the result is irrelevant, the mapped function only validates + @type_info.tree_map_type(result_collection_constructor=lambda _, elts: None) + def _validate_field_dims(type_: ts.TypeSpec) -> None: + # non-field leaves are invalid `out` types, but only rejected later, by the + # check of the `out` type against the operator return type + if isinstance(type_, ts.FieldType) and list(type_.dims) != dict_dims: + raise ValueError( + f"Dimensions in out field and field domain are not equivalent: " + f"expected '{type_.dims}', got '{dict_dims}'." + ) + + _validate_field_dims(out) + for domain_values in dom.values_: if len(domain_values.elts) != 2: raise ValueError( @@ -68,11 +83,18 @@ def _validate_domain_out( ) elif isinstance(dom, past.TupleExpr): + if isinstance(out, ts.VarArgType): + # The fixed-length domain tuple would pin the length of the variable-length + # output, so any call with a different length would break the program. + raise ValueError( + f"Per-element domain tuples are not supported for the variable-length tuple " + f"output '{out}'; use a single domain that applies to every element." + ) if isinstance(out, ts.CollectionTypeSpec): # TODO(havogt): for NamedCollectionType, we assume the domain tuple is in element order out_elts = out.types else: - raise ValueError(f"Tuple domain requires tuple output, got {type(out)}.") + raise ValueError(f"Tuple domain requires tuple output, got '{out}'.") if len(dom.elts) != len(out_elts): raise ValueError("Mismatched tuple lengths between domain and output.") @@ -142,6 +164,14 @@ def visit_Slice(self, node: past.Slice, **kwargs: Any) -> past.Slice: def visit_Subscript(self, node: past.Subscript, **kwargs: Any) -> past.Subscript: value = self.visit(node.value, **kwargs) + if isinstance(value.type, (*ts.COLLECTION_TYPE_SPECS, ts.VarArgType)): + # A subscript on a program argument is a domain slice; applying one to a + # (possibly variable-length) tuple is not supported, and element selection is + # not supported for any tuple argument. + raise errors.DSLError( + node.location, + f"Subscripting a tuple is not supported, got '{value.type}'.", + ) return past.Subscript( value=value, slice_=self.visit(node.slice_, **kwargs), diff --git a/src/gt4py/next/ffront/past_to_itir.py b/src/gt4py/next/ffront/past_to_itir.py index 05af5fa937..aee05f0bdd 100644 --- a/src/gt4py/next/ffront/past_to_itir.py +++ b/src/gt4py/next/ffront/past_to_itir.py @@ -9,7 +9,6 @@ from __future__ import annotations import dataclasses -import functools from typing import Any, Optional, Sequence, cast import devtools @@ -27,7 +26,7 @@ ) from gt4py.next.ffront.stages import ConcretePASTProgramDef from gt4py.next.iterator import ir as itir -from gt4py.next.iterator.ir_utils import ir_makers as im +from gt4py.next.iterator.ir_utils import ir_makers as im, misc as ir_misc from gt4py.next.iterator.transforms import remap_symbols, replace_get_domain_range_with_constants from gt4py.next.otf import arguments, stages, workflow from gt4py.next.type_system import type_info, type_specifications as ts @@ -93,8 +92,10 @@ def past_to_gtir(inp: ConcretePASTProgramDef) -> stages.CompilableProgramDef: for gt_callable in gt_callables: lowered_funcs.append(gt_callable.__gt_gtir__()) + past_node = _specialize_program_params(inp.data.past_node, inp.args.args, inp.args.kwargs) + itir_program = ProgramLowering.apply( - inp.data.past_node, function_definitions=lowered_funcs, grid_type=grid_type + past_node, function_definitions=lowered_funcs, grid_type=grid_type ) # TODO(tehrengruber): Put this in a dedicated transformation step. @@ -209,29 +210,50 @@ def _compute_field_slice(node: past.Subscript) -> list[past.Slice]: return out_field_slice_ -def _get_element_from_tuple_expr(node: past.Expr, path: tuple[int, ...]) -> past.Expr: - """Get element from a (nested) TupleExpr by following the given path. - - Pre-condition: `node` is a `past.TupleExpr` (if `path ! = ()`) - and `path` is a valid path through the nested tuple structure. - """ - return functools.reduce(lambda e, i: e.elts[i], path, node) # type: ignore[attr-defined] # see pre-condition - - -def _unwrap_tuple_expr(expr: past.Expr, path: tuple[int, ...]) -> tuple[past.Expr, Sequence[int]]: +def _unwrap_tuple_expr( + expr: past.Expr, path: tuple[int | type_info.VarArgElementKey, ...] +) -> tuple[past.Expr, Sequence[int | type_info.VarArgElementKey]]: """Unwrap (nested) TupleExpr by following the given path as long as possible. If a non-tuple expression is encountered, the current expression and the remaining path are returned. """ - path_remainder: Sequence[int] = path + path_remainder: Sequence[int | type_info.VarArgElementKey] = path while isinstance(expr, past.TupleExpr): idx, *path_remainder = path_remainder + # a `TupleExpr` always has a fixed-length tuple type, so the path element is an index + assert isinstance(idx, int) expr = expr.elts[idx] return expr, path_remainder +def _specialize_program_params( + past_node: past.Program, arg_types: Sequence[ts.TypeSpec], kwarg_types: dict[str, ts.TypeSpec] +) -> past.Program: + """ + Replace the program parameter types by the concrete argument types of the call. + + In particular, parameters declared with a generic type, e.g. a variable-length tuple, + only have a concrete type at the call. The arguments are expected to be canonicalized + against the signature already (see `transform_program_args`). + """ + param_types: dict[str, ts.TypeSpec] = { + **{str(param.id): arg_type for param, arg_type in zip(past_node.params, arg_types)}, + **kwarg_types, + } + new_params = [] + for param in past_node.params: + param_type = param_types[str(param.id)] + # incompatible argument types have already been rejected in `transform_program_args` + assert type_info.is_concretizable(param.type, to_type=param_type), ( + f"Argument of type '{param_type}' is not compatible with parameter " + f"'{param.id}' of type '{param.type}'." + ) + new_params.append(param.copy(update={"type": param_type})) + return past_node.copy(update={"params": new_params}) + + @dataclasses.dataclass class ProgramLowering( traits.PreserveLocationVisitor, traits.VisitorWithSymbolTableTrait, NodeTranslator @@ -352,17 +374,18 @@ def _visit_slice_bound( lowered_bound.location = slice_bound.location return lowered_bound - def _construct_itir_out_arg(self, node: past.Expr) -> itir.Expr: + def _construct_itir_out_arg(self, node: past.Expr, allow_slicing: bool) -> itir.Expr: if isinstance(node, past.Name): return itir.SymRef(id=node.id, location=node.location) elif isinstance(node, past.Subscript): - itir_node = self._construct_itir_out_arg(node.value) + assert allow_slicing + itir_node = self._construct_itir_out_arg(node.value, allow_slicing) itir_node.location = node.location return itir_node elif isinstance(node, past.TupleExpr): return itir.FunCall( fun=itir.SymRef(id="make_tuple"), - args=[self._construct_itir_out_arg(el) for el in node.elts], + args=[self._construct_itir_out_arg(el, allow_slicing) for el in node.elts], location=node.location, ) else: @@ -375,71 +398,48 @@ def _construct_itir_domain_arg( self, out_expr: itir.Expr, out_type: ts.FieldType, - node_domain: Optional[past.Expr], slices: Optional[list[past.Slice]] = None, ) -> itir.FunCall: - domain_args = [] + ranges: dict[common.Dimension, tuple[itir.Expr, itir.Expr]] = {} for dim_i, dim in enumerate(out_type.dims): + if dim.kind == common.DimensionKind.LOCAL: + raise ValueError(f"common.Dimension '{dim.value}' must not be local.") # an expression for the range of a dimension - dim_range = im.call("get_domain_range")( - out_expr, itir.AxisLiteral(value=dim.value, kind=dim.kind) - ) + dim_range = im.call("get_domain_range")(out_expr, im.axis_literal(dim)) dim_start, dim_stop = im.tuple_get(0, dim_range), im.tuple_get(1, dim_range) - # bounds - lower: itir.Expr - upper: itir.Expr - if node_domain is not None: - assert isinstance(node_domain, past.Dict) - lower, upper = self._construct_itir_initialized_domain_arg(dim_i, dim, node_domain) - else: - lower = self._visit_slice_bound( - slices[dim_i].lower if slices else None, - dim_start, - dim_start, - dim_stop, - ) - upper = self._visit_slice_bound( - slices[dim_i].upper if slices else None, - dim_stop, - dim_start, - dim_stop, - ) - - if dim.kind == common.DimensionKind.LOCAL: - raise ValueError(f"common.Dimension '{dim.value}' must not be local.") - domain_args.append( - itir.FunCall( - fun=itir.SymRef(id="named_range"), - args=[itir.AxisLiteral(value=dim.value, kind=dim.kind), lower, upper], - ) + lower = self._visit_slice_bound( + slices[dim_i].lower if slices else None, + dim_start, + dim_start, + dim_stop, ) - - if self.grid_type == common.GridType.CARTESIAN: - domain_builtin = "cartesian_domain" - elif self.grid_type == common.GridType.UNSTRUCTURED: - domain_builtin = "unstructured_domain" - else: - raise AssertionError() - - return itir.FunCall( - fun=itir.SymRef(id=domain_builtin), - args=domain_args, - location=(node_domain or out_expr).location, - ) - - def _construct_itir_initialized_domain_arg( - self, dim_i: int, dim: common.Dimension, node_domain: past.Dict - ) -> list[itir.FunCall]: - assert len(node_domain.values_[dim_i].elts) == 2 - keys_dims_types = cast(ts.DimensionType, node_domain.keys_[dim_i].type).dim - if keys_dims_types != dim: - raise ValueError( - "common.Dimensions in out field and field domain are not equivalent:" - f"expected '{dim}', got '{keys_dims_types}'." + upper = self._visit_slice_bound( + slices[dim_i].upper if slices else None, + dim_stop, + dim_start, + dim_stop, ) - - return [self.visit(bound) for bound in node_domain.values_[dim_i].elts] + ranges[dim] = (lower, upper) + + domain_expr = im.domain(self.grid_type, ranges) + domain_expr.location = out_expr.location + return domain_expr + + def _lower_domain_arg(self, domain_arg: past.Expr) -> itir.Expr: + if isinstance(domain_arg, past.TupleExpr): + return im.make_tuple(*(self._lower_domain_arg(el) for el in domain_arg.elts)) + assert isinstance(domain_arg, past.Dict) + ranges = { + cast(ts.DimensionType, dim_name.type).dim: ( + self.visit(bounds.elts[0]), + self.visit(bounds.elts[1]), + ) + for dim_name, bounds in zip(domain_arg.keys_, domain_arg.values_, strict=True) + } + domain_expr = im.domain(self.grid_type, ranges) + domain_expr.location = domain_arg.location + return domain_expr def _split_field_and_slice( self, field: past.Name | past.Subscript @@ -452,40 +452,43 @@ def _split_field_and_slice( def _visit_stencil_call_out_arg( self, out_arg: past.Expr, domain_arg: Optional[past.Expr], **kwargs: Any - ) -> tuple[itir.Expr, itir.FunCall]: + ) -> tuple[itir.Expr, itir.Expr]: assert isinstance(out_arg, (past.Subscript, past.Name, past.TupleExpr)), ( "Unexpected 'out' argument. Must be a 'past.Subscript', 'past.Name' or 'past.TupleExpr' node." ) + assert out_arg.type is not None + + if domain_arg is not None: + # Mutual exclusivity of a `domain` argument and slices is ensured in the type + # deduction, so the target is just the plain lowered `out` argument. Note that + # a single domain may apply to an entire tuple subtree of the target (see + # `itir.SetAt`). + return self._construct_itir_out_arg( + out_arg, allow_slicing=False + ), self._lower_domain_arg(domain_arg) + else: - @type_info.tree_map_type( - with_path_arg=True, - unpack=True, - result_collection_constructor=lambda _, elts: im.make_tuple(*elts), - ) - def impl(out_type: ts.FieldType, path: tuple[int, ...]) -> tuple[itir.Expr, itir.Expr]: - out_field, path_remainder = _unwrap_tuple_expr(out_arg, path) - - assert isinstance(out_field, (past.Name, past.Subscript)) - out_field, slice_info = self._split_field_and_slice(out_field) + @ir_misc.tree_map_type_into_expr(with_path_arg=True) + def impl( + out_type: ts.FieldType, + current: itir.Expr, + path: tuple[int | type_info.VarArgElementKey, ...], + ) -> itir.Expr: + out_field, path_remainder = _unwrap_tuple_expr(out_arg, path) + + # A subscript (domain slice) can only appear on a single field, i.e. on an + # expression that does not decompose into further tuple elements. + assert isinstance(out_field, past.Name) or ( + isinstance(out_field, past.Subscript) and len(path_remainder) == 0 + ) + _, slice_info = self._split_field_and_slice(out_field) - domain_element = ( - _get_element_from_tuple_expr(domain_arg, path) - if isinstance(domain_arg, past.TupleExpr) - else domain_arg - ) + return self._construct_itir_domain_arg(current, out_type, slice_info) - lowered_out_field = functools.reduce( - lambda expr, i: im.tuple_get(i, expr), path_remainder, self.visit(out_field) - ) - lowered_domain = self._construct_itir_domain_arg( - lowered_out_field, - out_type, - domain_element, - slice_info, - ) - return lowered_out_field, lowered_domain + lowered_out = self._construct_itir_out_arg(out_arg, allow_slicing=True) + lowered_domain = impl(out_arg.type, lowered_out) - return impl(out_arg.type) + return lowered_out, lowered_domain def visit_Constant(self, node: past.Constant, **kwargs: Any) -> itir.Literal: if isinstance(node.type, ts.ScalarType) and node.type.shape is None: diff --git a/src/gt4py/next/ffront/type_info.py b/src/gt4py/next/ffront/type_info.py index e7280d51b2..8eb3a516ae 100644 --- a/src/gt4py/next/ffront/type_info.py +++ b/src/gt4py/next/ffront/type_info.py @@ -153,9 +153,9 @@ def function_signature_incompatibilities_fieldop( def _tree_map_type_constructor_drop_python_type( - value: ts.CollectionTypeSpecT, + value: ts.CollectionTypeSpecT | ts.VarArgType, elems: Iterable[ts.DataType | ts.DimensionType | ts.DeferredType], -) -> ts.CollectionTypeSpecT: +) -> ts.CollectionTypeSpecT | ts.VarArgType: result = type_info.tree_map_type_constructor(value, elems) if isinstance(value, ts.NamedCollectionType): result = datamodels.evolve(result, original_python_type=ts.ANY_PYTHON_TYPE_NAME) diff --git a/src/gt4py/next/iterator/ir.py b/src/gt4py/next/iterator/ir.py index f024ef6168..6a237ffd39 100644 --- a/src/gt4py/next/iterator/ir.py +++ b/src/gt4py/next/iterator/ir.py @@ -125,6 +125,16 @@ class Stmt(Node): ... class SetAt(Stmt): # from JAX array.at[...].set() + """ + Write `expr`, computed on `domain`, into `target`. + + For a tuple-valued `target` the `domain` is a matching tuple of domains (built with + `make_tuple`), except that a single domain may also apply to an entire tuple subtree + of the target, i.e. to each of its elements. This shallower structure is needed for + variable-length tuple outputs, whose concrete length is unknown when the domain + expression is constructed. + """ + expr: Expr # only `as_fieldop(stencil)(inp0, ...)` in first refactoring domain: Expr target: Expr # `make_tuple` or SymRef @@ -153,6 +163,13 @@ class Program(Node, ValidatedSymbolTableTrait): Sym(id=name) for name in sorted(BUILTINS) ] # sorted for serialization stability + def __hash__(self) -> int: + # Only here do types contribute to the hash as they are the starting point for + # type inference, whereas in all other places the types can be recovered. Note + # that equality still ignores types, so equal programs may have different hashes; + # the users only consume the hash value itself, never using programs as dict keys. + return hash((Node.__hash__(self), tuple(str(param.type) for param in self.params))) + # TODO(fthaler): just use hashable types in nodes (tuples instead of lists) Sym.__hash__ = Node.__hash__ # type: ignore[method-assign] @@ -165,7 +182,6 @@ class Program(Node, ValidatedSymbolTableTrait): Lambda.__hash__ = Node.__hash__ # type: ignore[method-assign] FunCall.__hash__ = Node.__hash__ # type: ignore[method-assign] FunctionDefinition.__hash__ = Node.__hash__ # type: ignore[method-assign] -Program.__hash__ = Node.__hash__ # type: ignore[method-assign] SetAt.__hash__ = Node.__hash__ # type: ignore[method-assign] IfStmt.__hash__ = Node.__hash__ # type: ignore[method-assign] InfinityLiteral.__hash__ = Node.__hash__ # type: ignore[method-assign] diff --git a/src/gt4py/next/iterator/ir_utils/misc.py b/src/gt4py/next/iterator/ir_utils/misc.py index d9a6e85f57..3ae0af8bf8 100644 --- a/src/gt4py/next/iterator/ir_utils/misc.py +++ b/src/gt4py/next/iterator/ir_utils/misc.py @@ -7,8 +7,9 @@ # SPDX-License-Identifier: BSD-3-Clause import dataclasses +import functools from collections import ChainMap -from typing import Callable, Iterable, TypeVar, cast +from typing import Any, Callable, Iterable, Optional, TypeVar, cast from gt4py import eve from gt4py._core import definitions as core_defs @@ -17,7 +18,7 @@ from gt4py.next.iterator import embedded, ir as itir from gt4py.next.iterator.ir_utils import common_pattern_matcher as cpm, ir_makers as im from gt4py.next.iterator.transforms import inline_lambdas -from gt4py.next.type_system import type_specifications as ts +from gt4py.next.type_system import type_info, type_specifications as ts @dataclasses.dataclass(frozen=True) @@ -279,3 +280,86 @@ def value_from_literal(literal: itir.Literal) -> core_defs.Scalar: assert literal.value in ["True", "False"] return literal.value == "True" return getattr(embedded, str(literal.type))(literal.value) + + +#: Name bound to the value of the current element in `tree_map_type_into_expr`; only the +#: (shadowed) binders introduced by `map_tuple` lambdas remain visible in the result. +_CURRENT_ELEMENT = "__tmtie_el" + + +def tree_map_type_into_expr( + fun: Optional[Callable[..., itir.Expr]] = None, + *, + with_path_arg: bool = False, + unpack: bool = False, +) -> Callable[..., Any]: + """ + Map `fun` over the leaves of a type tree, building an ITIR expression over a + corresponding value expression. + + Like `type_info.tree_map_type`, but `fun` additionally receives a `current` argument — + a reference to the value of the current element — and the mapped callable is called + with the value expression alongside the type. Fixed-length collection nodes recombine + the element results with `make_tuple`, rebinding `current` per element; `VarArgType` + nodes map their elements with the `map_tuple` builtin, so the result is also valid for + types whose concrete length is only known later. The bindings introduced along the way + are inlined unconditionally before returning — the value expression should hence be + cheap (e.g. a symbol reference or tuple accesses thereon), as it may be duplicated. + With `with_path_arg`, paths contain `type_info.VAR_ARG_ELEMENT` below variable-length + tuples. + + Can be used directly or as a (possibly parametrized) decorator. + + >>> field = ts.FieldType(dims=[], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64)) + >>> type_ = ts.TupleType(types=[field, ts.VarArgType(element_type=field)]) + >>> expr = tree_map_type_into_expr(lambda t, current: im.call("get_dom")(current))( + ... type_, im.ref("out") + ... ) + >>> str(expr) + '{get_dom(out[0]), map_tuple(λ(__tmtie_el_) → get_dom(__tmtie_el_))(out[1])}' + """ + if fun is None: + return functools.partial( + tree_map_type_into_expr, with_path_arg=with_path_arg, unpack=unpack + ) + + @functools.wraps(fun) + def impl(type_: ts.TypeSpec, expr: itir.Expr) -> itir.Expr | tuple[itir.Expr, ...]: + current = im.ref(_CURRENT_ELEMENT) + + def collection_constructor(value: ts.TypeSpec, elts: Iterable[itir.Expr]) -> itir.Expr: + if isinstance(value, ts.VarArgType): + (element_expr,) = elts + return im.call(im.call("map_tuple")(im.lambda_(_CURRENT_ELEMENT)(element_expr)))( + current + ) + return im.make_tuple( + *( + im.let(_CURRENT_ELEMENT, im.tuple_get(i, current))(el) + for i, el in enumerate(elts) + ) + ) + + leaf_fun = ( + (lambda type_, path: fun(type_, current, path)) + if with_path_arg + else (lambda type_: fun(type_, current)) + ) + mapped = type_info.tree_map_type( + leaf_fun, + with_path_arg=with_path_arg, + unpack=unpack, + result_collection_constructor=collection_constructor, + ) + + def bind(tree: itir.Expr) -> itir.Expr: + # inline unconditionally so consumers see plain tuple accesses instead of + # `let`-bound references (see docstring) + return inline_lambdas.InlineLambdas.apply(im.let(_CURRENT_ELEMENT, expr)(tree)) + + result = mapped(type_) + if unpack: + return tuple(bind(tree) for tree in cast(tuple[itir.Expr, ...], result)) + return bind(result) + + return impl diff --git a/src/gt4py/next/otf/compiled_program.py b/src/gt4py/next/otf/compiled_program.py index 4e0cc75dc5..99d9ac1edf 100644 --- a/src/gt4py/next/otf/compiled_program.py +++ b/src/gt4py/next/otf/compiled_program.py @@ -733,6 +733,12 @@ def compile( pool.compile(static_arg0=[0], static_arg1=[2]).compile(static_arg=[1], static_arg1=[3]) will compile for (0,2), (1,3) """ + if self._is_generic: + raise ValueError( + f"Cannot compile '{self.definition_stage.definition.__name__}' ahead of time: " + "generic programs (e.g. with variable-length tuple parameters or direct scan " + "operator calls) require the concrete argument types of a call." + ) for offset_provider in offset_providers: # not included in product for better type checking for static_values in itertools.product(*static_args.values()): # The argument descriptors built here must be kept in sync with the diff --git a/src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg.py b/src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg.py index e645068f64..2442939bbf 100644 --- a/src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg.py +++ b/src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg.py @@ -1151,6 +1151,10 @@ def _visit_target( # Visit the field operator expression. source_tree = self._visit_expression(stmt.expr, ctx) + + # A single domain may apply to an entire tuple subtree of the target (see + # `gtir.SetAt`); promote it so all trees have the same structure. + _, domain = gtx_utils.equalize_tuple_structure(source_tree, domain, bidirectional=False) ctx_input_data = ctx.get_input_data() # In order to support inout argument, write the result in separate next state diff --git a/src/gt4py/next/type_system/type_info.py b/src/gt4py/next/type_system/type_info.py index 2600746ee5..023bb080f8 100644 --- a/src/gt4py/next/type_system/type_info.py +++ b/src/gt4py/next/type_system/type_info.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import functools +import itertools import types from collections.abc import Callable, Collection, Iterable, Iterator from typing import Any, Final, Literal, Sequence, Type, TypeGuard, TypeVar, overload @@ -83,6 +84,19 @@ def type_class(symbol_type: ts.TypeSpec) -> Type[ts.TypeSpec]: ) +class VarArgElementKey: + """Path key of the representative element of a `VarArgType` in type traversals.""" + + def __repr__(self) -> str: + return "VAR_ARG_ELEMENT" + + +#: Type traversals (`tree_map_type`, `primitive_constituents`) visit a `VarArgType` as +#: a single representative element (its `element_type`), addressed by this sentinel +#: instead of a position. +VAR_ARG_ELEMENT: Final = VarArgElementKey() + + @overload def primitive_constituents( symbol_type: ts.TypeSpec, with_path_arg: Literal[False] = False @@ -92,15 +106,21 @@ def primitive_constituents( @overload def primitive_constituents( symbol_type: ts.TypeSpec, with_path_arg: Literal[True] -) -> utils.XIterable[tuple[ts.TypeSpec, tuple[int, ...]]]: ... +) -> utils.XIterable[tuple[ts.TypeSpec, tuple[int | VarArgElementKey, ...]]]: ... def primitive_constituents( symbol_type: ts.TypeSpec, with_path_arg: bool = False -) -> utils.XIterable[ts.TypeSpec] | utils.XIterable[tuple[ts.TypeSpec, tuple[int, ...]]]: +) -> ( + utils.XIterable[ts.TypeSpec] + | utils.XIterable[tuple[ts.TypeSpec, tuple[int | VarArgElementKey, ...]]] +): """ Return the primitive types contained in a composite type. + A `VarArgType` is traversed as its single representative element, addressed by + `VAR_ARG_ELEMENT` in paths. + >>> from gt4py.next import common >>> I = common.Dimension(value="I") >>> int_type = ts.ScalarType(kind=ts.ScalarKind.INT64) @@ -116,9 +136,11 @@ def primitive_constituents( """ def constituents_yielder( - symbol_type: ts.TypeSpec, path: tuple[int, ...] - ) -> Iterator[ts.TypeSpec] | Iterator[tuple[ts.TypeSpec, tuple[int, ...]]]: - if isinstance(symbol_type, ts.COLLECTION_TYPE_SPECS): + symbol_type: ts.TypeSpec, path: tuple[int | VarArgElementKey, ...] + ) -> Iterator[ts.TypeSpec] | Iterator[tuple[ts.TypeSpec, tuple[int | VarArgElementKey, ...]]]: + if isinstance(symbol_type, ts.VarArgType): + yield from constituents_yielder(symbol_type.element_type, (*path, VAR_ARG_ELEMENT)) + elif isinstance(symbol_type, ts.COLLECTION_TYPE_SPECS): for i, el_type in enumerate(symbol_type.types): yield from constituents_yielder(el_type, (*path, i)) else: @@ -134,10 +156,27 @@ def constituents_yielder( _C = TypeVar("_C") +def _tree_map_type_collection_elements(value: ts.TypeSpec) -> Iterable[ts.TypeSpec]: + if isinstance(value, ts.VarArgType): + return (value.element_type,) + assert isinstance(value, ts.COLLECTION_TYPE_SPECS) + return value.types + + +def _tree_map_type_collection_keys(value: ts.TypeSpec) -> Iterable[Any]: + if isinstance(value, ts.VarArgType): + return (VAR_ARG_ELEMENT,) + return itertools.count() + + def tree_map_type_constructor( - value: ts.CollectionTypeSpecT, + value: ts.CollectionTypeSpecT | ts.VarArgType, elems: Iterable[ts.DataType | ts.DimensionType | ts.DeferredType], -) -> ts.CollectionTypeSpecT: +) -> ts.CollectionTypeSpecT | ts.VarArgType: + if isinstance(value, ts.VarArgType): + (element_type,) = elems + assert isinstance(element_type, ts.DataType) + return ts.VarArgType(element_type=element_type) return ( ts.NamedCollectionType( keys=value.keys, original_python_type=value.original_python_type, types=list(elems) @@ -194,8 +233,9 @@ def tree_map_type( ) return next_utils.tree_map( fun, - collection_type=ts.COLLECTION_TYPE_SPECS, - collection_elements=lambda t: t.types, + collection_type=(*ts.COLLECTION_TYPE_SPECS, ts.VarArgType), + collection_elements=_tree_map_type_collection_elements, + collection_keys=_tree_map_type_collection_keys, result_collection_constructor=result_collection_constructor, with_path_arg=with_path_arg, unpack=unpack, diff --git a/src/gt4py/next/type_system/type_specifications.py b/src/gt4py/next/type_system/type_specifications.py index 7c9e089096..73a8a08a5d 100644 --- a/src/gt4py/next/type_system/type_specifications.py +++ b/src/gt4py/next/type_system/type_specifications.py @@ -195,6 +195,8 @@ def __str__(self) -> str: return f"NamedTuple{{{', '.join(f'{k}: {v}' for k, v in zip(self.keys, self.types))}}}" +# TODO(tehrengruber): Add VarArgType to this list, which requires checking all users of it +# that potentially depend on the length being known. CollectionTypeSpecT = TypeVar("CollectionTypeSpecT", TupleType, NamedCollectionType) CollectionTypeSpec = TupleType | NamedCollectionType COLLECTION_TYPE_SPECS: Final[tuple[type[CollectionTypeSpec], ...]] = typing.get_args( diff --git a/src/gt4py/next/utils.py b/src/gt4py/next/utils.py index 54491d3663..7bf6ec5a5f 100644 --- a/src/gt4py/next/utils.py +++ b/src/gt4py/next/utils.py @@ -26,7 +26,7 @@ overload, ) -from gt4py.eve import utils as eve_utils +from gt4py.eve import NOTHING, utils as eve_utils _T = TypeVar("_T") @@ -153,6 +153,9 @@ def tree_map( reconstruct a collection, since `result_collection_constructor` receives the original collection and can recover them from there. result_collection_constructor: Type of the collection to be returned. If `None` the same type as `collection_type` is used. + `fun` is guaranteed to have been applied to every element before the constructor + is called (the mapped elements are materialized eagerly), so a constructor may + discard the results, e.g. for validation-only maps. unpack: Replicate tuple structure returned from `fun` to the mapped result, i.e. return tuple of result collections instead of result collections of tuples. with_path_arg: Pass the path to access the current element to `fun`. @@ -298,13 +301,14 @@ def impl(*args: Any | tuple[Any | tuple, ...]) -> _R | tuple[_R | tuple, ...]: def equalize_tuple_structure( - d1: _T1, d2: _T2, *, fill_value: Any, bidirectional: bool = True + d1: _T1, d2: _T2, *, fill_value: Any = NOTHING, bidirectional: bool = True ) -> tuple[_T1, _T2]: """ Given two values or tuples thereof equalize their structure. If one of the arguments is a tuple the argument will be promoted to a tuple of the same - structure. + structure. By default only the structure is equalized and differing tuple lengths are + an error; pass `fill_value` to instead pad the shorter tuple with it. >>> equalize_tuple_structure((1,), (2, 3), fill_value=None) == ( ... (1, None), @@ -342,6 +346,8 @@ def equalize_tuple_structure( if d1_is_tuple and d2_is_tuple: if not bidirectional and len(d1) < len(d2): # type: ignore[arg-type] # ensured by condition above raise ValueError(f"Expected `{d2!s}` to be longer than or equal to `{d1!s}`.") + if fill_value is NOTHING and len(d1) != len(d2): # type: ignore[arg-type] # ensured by condition above + raise ValueError(f"Expected `{d1!s}` and `{d2!s}` to have equal lengths.") return tuple( # type: ignore[return-value] # mypy not smart enough zip( *( diff --git a/tests/next_tests/integration_tests/cases.py b/tests/next_tests/integration_tests/cases.py index 9ba5f65077..8577243008 100644 --- a/tests/next_tests/integration_tests/cases.py +++ b/tests/next_tests/integration_tests/cases.py @@ -24,7 +24,7 @@ import inspect import types import typing -from typing import Any, Callable, Literal, Optional, Mapping, Protocol, Self, TypeAlias +from typing import Any, Callable, Final, Literal, Optional, Mapping, Protocol, Self, TypeAlias import numpy as np import pytest @@ -347,6 +347,7 @@ def allocate( strategy: Optional[DataInitializer] = None, dtype: Optional[np.typing.DTypeLike] = None, extend: Optional[dict[gtx.Dimension, tuple[int, int]]] = None, + varargs_length: Optional[int] = None, ) -> FieldViewArg: """ Allocate a parameter or return value from a fieldview code object. @@ -364,6 +365,8 @@ def allocate( extend: Lower and upper size extension per dimension. Useful for shifted fields, which must start off bigger than the output field in the shifted dimension. + varargs_length: Concrete length to allocate for variable-length tuple + parameters (default: ``DEFAULT_VARARGS_LENGTH``). """ if sizes: assert not domain and all(dim in case.default_sizes for dim in sizes) @@ -393,6 +396,7 @@ def allocate( domain=domain, dtype=dtype, strategy=strategy.from_case(case=case, fieldview_prog=fieldview_prog, arg_name=name), + varargs_length=varargs_length, ) @@ -588,6 +592,11 @@ def unstructured_case_3d(unstructured_case): ) +#: Default concrete length used for variable-length tuple parameters, e.g. when +#: allocating arguments for them. +DEFAULT_VARARGS_LENGTH: Final = 3 + + def _allocate_from_type( case: Case, arg_type: ts.TypeSpec, @@ -595,6 +604,7 @@ def _allocate_from_type( strategy: DataInitializer, dtype: Optional[np.typing.DTypeLike] = None, tuple_start: Optional[int] = None, + varargs_length: Optional[int] = None, ) -> FieldViewArg: """Allocate data based on the type or a (nested) tuple thereof.""" match arg_type: @@ -610,18 +620,32 @@ def _allocate_from_type( return tuple( ( _allocate_from_type( - case=case, arg_type=t, domain=domain, dtype=dtype, strategy=strategy + case=case, + arg_type=t, + domain=domain, + dtype=dtype, + strategy=strategy, + varargs_length=varargs_length, ) for t in types ) ) case ts.VarArgType(element_type=element_type): + length = varargs_length if varargs_length is not None else DEFAULT_VARARGS_LENGTH + # the unique-value offsets are computed from the declared types using the + # default length (see `get_param_size`) + assert length <= DEFAULT_VARARGS_LENGTH return tuple( ( _allocate_from_type( - case=case, arg_type=t, domain=domain, dtype=dtype, strategy=strategy + case=case, + arg_type=t, + domain=domain, + dtype=dtype, + strategy=strategy, + varargs_length=varargs_length, ) - for t in [element_type] * 3 # TODO: revisit + for t in [element_type] * length ) ) case ts.NamedCollectionType(types=types) as named_collection_type_spec: @@ -634,7 +658,12 @@ def _allocate_from_type( tuple( ( _allocate_from_type( - case=case, arg_type=t, domain=domain, dtype=dtype, strategy=strategy + case=case, + arg_type=t, + domain=domain, + dtype=dtype, + strategy=strategy, + varargs_length=varargs_length, ) for t in types ) @@ -670,7 +699,9 @@ def get_param_size(param_type: ts.TypeSpec, sizes: dict[gtx.Dimension, int]) -> case ts.NamedCollectionType(types=types): return sum([get_param_size(t, sizes=sizes) for t in types]) case ts.VarArgType(element_type=element_type): - return get_param_size(ts.TupleType(types=[element_type] * 3), sizes) # TODO: revisit + return get_param_size( + ts.TupleType(types=[element_type] * DEFAULT_VARARGS_LENGTH), sizes + ) case _: raise TypeError(f"Can not get size for parameter of type '{param_type}'.") diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_program.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_program.py index d882a88ec4..62d80a8e47 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_program.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_program.py @@ -250,17 +250,13 @@ def test_dimensions_domain(cartesian_case): def empty_domain_fieldop(a: cases.IJField): return a - @gtx.program - def empty_domain_program(a: cases.IJField, out_field: cases.IJField): - empty_domain_fieldop(a, out=out_field, domain={JDim: (0, 1), IDim: (0, 1)}) - - a = cases.allocate(cartesian_case, empty_domain_program, "a")() - out_field = cases.allocate(cartesian_case, empty_domain_program, "out_field")() - with pytest.raises( - ValueError, match=(r"Dimensions in out field and field domain are not equivalent") + errors.DSLError, match=(r"Dimensions in out field and field domain are not equivalent") ): - cases.run(cartesian_case, empty_domain_program, a, out_field, offset_provider={}) + + @gtx.program + def empty_domain_program(a: cases.IJField, out_field: cases.IJField): + empty_domain_fieldop(a, out=out_field, domain={JDim: (0, 1), IDim: (0, 1)}) @pytest.mark.uses_origin diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py index 2b18039e7c..63983ead54 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py @@ -189,8 +189,8 @@ def testee(tracers: tuple[cases.IField, ...], factor: int32) -> tuple[cases.IFie return tuple(tracer * factor for tracer in tracers) for length in (3, 2): - tracers = cases.allocate(cartesian_case, testee, "tracers")()[:length] - out = cases.allocate(cartesian_case, testee, cases.RETURN)()[:length] + tracers = cases.allocate(cartesian_case, testee, "tracers", varargs_length=length)() + out = cases.allocate(cartesian_case, testee, cases.RETURN, varargs_length=length)() cases.verify( cartesian_case, testee, @@ -201,6 +201,148 @@ def testee(tracers: tuple[cases.IField, ...], factor: int32) -> tuple[cases.IFie ) +@pytest.mark.uses_tuple_args +def test_var_len_tuple_program(cartesian_case): + @gtx.field_operator + def scale(tracers: tuple[cases.IField, ...], factor: int32) -> tuple[cases.IField, ...]: + return tuple(tracer * factor for tracer in tracers) + + @gtx.program + def testee(tracers: tuple[cases.IField, ...], factor: int32, out: tuple[cases.IField, ...]): + scale(tracers, factor, out=out) + + for length in (3, 2): # ensure each concrete length gets its own specialization + tracers = cases.allocate(cartesian_case, testee, "tracers", varargs_length=length)() + out = cases.allocate(cartesian_case, testee, "out", varargs_length=length)() + cases.run(cartesian_case, testee, tracers, 42, out) + assert all(np.allclose(o.asnumpy(), t.asnumpy() * 42) for o, t in zip(out, tracers)) + + +@pytest.mark.uses_tuple_args +def test_var_len_tuple_program_domain(cartesian_case): + # A single `domain` argument applies to (is broadcast over) every element of the + # variable-length tuple output. + @gtx.field_operator + def scale(tracers: tuple[cases.IField, ...], factor: int32) -> tuple[cases.IField, ...]: + return tuple(tracer * factor for tracer in tracers) + + @gtx.program + def testee( + tracers: tuple[cases.IField, ...], + factor: int32, + size: int32, + out: tuple[cases.IField, ...], + ): + scale(tracers, factor, out=out, domain={IDim: (1, size - 1)}) + + size = cartesian_case.default_sizes[IDim] + tracers = cases.allocate(cartesian_case, testee, "tracers")() + out = cases.allocate(cartesian_case, testee, "out")() + cases.run(cartesian_case, testee, tracers, 42, int32(size), out) + for o, t in zip(out, tracers): + assert np.allclose(o.asnumpy()[1:-1], t.asnumpy()[1:-1] * 42) + assert np.all(o.asnumpy()[[0, -1]] == 0) + + +@pytest.mark.uses_tuple_args +def test_var_len_tuple_program_fixed_in_vararg_out(cartesian_case): + @gtx.field_operator + def scale_pairs( + pairs: tuple[tuple[cases.IField, cases.IField], ...], factor: int32 + ) -> tuple[tuple[cases.IField, cases.IField], ...]: + return tuple((a * factor, b * factor) for a, b in pairs) + + @gtx.program + def testee( + pairs: tuple[tuple[cases.IField, cases.IField], ...], + factor: int32, + out: tuple[tuple[cases.IField, cases.IField], ...], + ): + scale_pairs(pairs, factor, out=out) + + pairs = cases.allocate(cartesian_case, testee, "pairs")() + out = cases.allocate(cartesian_case, testee, "out")() + cases.run(cartesian_case, testee, pairs, 42, out) + assert all( + np.allclose(o.asnumpy(), p.asnumpy() * 42) + for o_pair, p_pair in zip(out, pairs) + for o, p in zip(o_pair, p_pair) + ) + + +@pytest.mark.uses_tuple_args +@pytest.mark.uses_tuple_returns +def test_var_len_tuple_program_vararg_in_fixed_out(cartesian_case): + @gtx.field_operator + def scale_and_copy( + tracers: tuple[cases.IField, ...], rho: cases.IField, factor: int32 + ) -> tuple[tuple[cases.IField, ...], cases.IField]: + return tuple(tracer * factor for tracer in tracers), rho + + @gtx.program + def testee( + tracers: tuple[cases.IField, ...], + rho: cases.IField, + factor: int32, + out: tuple[tuple[cases.IField, ...], cases.IField], + ): + scale_and_copy(tracers, rho, factor, out=out) + + tracers = cases.allocate(cartesian_case, testee, "tracers")() + rho = cases.allocate(cartesian_case, testee, "rho")() + out = cases.allocate(cartesian_case, testee, "out")() + cases.run(cartesian_case, testee, tracers, rho, 42, out) + assert all(np.allclose(o.asnumpy(), t.asnumpy() * 42) for o, t in zip(out[0], tracers)) + assert np.allclose(out[1].asnumpy(), rho.asnumpy()) + + +@pytest.mark.uses_tuple_args +@pytest.mark.uses_tuple_returns +def test_var_len_tuple_program_composed_out(cartesian_case): + # The 'out' argument is composed from several program parameters inside the program. + @gtx.field_operator + def scale_and_copy( + tracers: tuple[cases.IField, ...], rho: cases.IField, factor: int32 + ) -> tuple[tuple[cases.IField, ...], cases.IField]: + return tuple(tracer * factor for tracer in tracers), rho + + @gtx.program + def testee( + tracers: tuple[cases.IField, ...], + rho: cases.IField, + factor: int32, + out_tracers: tuple[cases.IField, ...], + out_rho: cases.IField, + ): + scale_and_copy(tracers, rho, factor, out=(out_tracers, out_rho)) + + tracers = cases.allocate(cartesian_case, testee, "tracers")() + rho = cases.allocate(cartesian_case, testee, "rho")() + out_tracers = cases.allocate( + cartesian_case, testee, "out_tracers", strategy=cases.ZeroInitializer() + )() + out_rho = cases.allocate(cartesian_case, testee, "out_rho", strategy=cases.ZeroInitializer())() + cases.run(cartesian_case, testee, tracers, rho, 42, out_tracers, out_rho) + assert all(np.allclose(o.asnumpy(), t.asnumpy() * 42) for o, t in zip(out_tracers, tracers)) + assert np.allclose(out_rho.asnumpy(), rho.asnumpy()) + + +@pytest.mark.uses_tuple_args +def test_var_len_tuple_program_compile_rejected(cartesian_case): + @gtx.field_operator + def scale(tracers: tuple[cases.IField, ...], factor: int32) -> tuple[cases.IField, ...]: + return tuple(tracer * factor for tracer in tracers) + + @gtx.program(backend=cartesian_case.backend) + def testee(tracers: tuple[cases.IField, ...], factor: int32, out: tuple[cases.IField, ...]): + scale(tracers, factor, out=out) + + if cartesian_case.backend is None: + pytest.skip("Ahead-of-time compilation requires a backend.") + with pytest.raises(ValueError, match="generic programs"): + testee.compile(offset_provider={}) + + @pytest.mark.uses_tuple_args @pytest.mark.xfail( strict=True, @@ -218,9 +360,9 @@ def testee( ) -> tuple[cases.IField, ...]: return a if c else b - a = cases.allocate(cartesian_case, testee, "a")()[:2] - b = cases.allocate(cartesian_case, testee, "b")() # allocated length: 3 - out = cases.allocate(cartesian_case, testee, cases.RETURN)()[:2] + a = cases.allocate(cartesian_case, testee, "a", varargs_length=2)() + b = cases.allocate(cartesian_case, testee, "b", varargs_length=3)() + out = cases.allocate(cartesian_case, testee, cases.RETURN, varargs_length=2)() with pytest.raises(errors.DSLError): cases.run(cartesian_case, testee, a, b, False, out=out) diff --git a/tests/next_tests/unit_tests/ffront_tests/test_func_to_past.py b/tests/next_tests/unit_tests/ffront_tests/test_func_to_past.py index 8486885076..7cc86a5a0b 100644 --- a/tests/next_tests/unit_tests/ffront_tests/test_func_to_past.py +++ b/tests/next_tests/unit_tests/ffront_tests/test_func_to_past.py @@ -152,6 +152,82 @@ def test_copy_restrict_parsing(copy_restrict_program_def): pattern_node.match(past_node, raise_exception=True) +@pytest.fixture +def vararg_scale_op(): + @gtx.field_operator + def scale( + tracers: tuple[gtx.Field[[IDim], float64], ...], factor: float64 + ) -> tuple[gtx.Field[[IDim], float64], ...]: + return tuple(tracer * factor for tracer in tracers) + + return scale + + +def test_vararg_subscript_out_exception(vararg_scale_op): + scale = vararg_scale_op + + def vararg_slice_program( + tracers: tuple[gtx.Field[[IDim], float64], ...], + factor: float64, + out: tuple[gtx.Field[[IDim], float64], ...], + ): + scale(tracers, factor, out=out[0:2]) + + with pytest.raises(errors.DSLError, match="Subscripting a tuple"): + ProgramParser.apply_to_function(vararg_slice_program) + + +def test_vararg_subscript_arg_exception(vararg_scale_op): + scale = vararg_scale_op + + def vararg_slice_program( + tracers: tuple[gtx.Field[[IDim], float64], ...], + factor: float64, + out: tuple[gtx.Field[[IDim], float64], ...], + ): + scale(tracers[0:2], factor, out=out) + + with pytest.raises(errors.DSLError, match="Subscripting a tuple"): + ProgramParser.apply_to_function(vararg_slice_program) + + +def test_tuple_subscript_out_exception(): + @gtx.field_operator + def pair( + a: gtx.Field[[IDim], float64], + ) -> tuple[gtx.Field[[IDim], float64], gtx.Field[[IDim], float64]]: + return a, a + + def tuple_slice_program( + a: gtx.Field[[IDim], float64], + out: tuple[gtx.Field[[IDim], float64], gtx.Field[[IDim], float64]], + ): + pair(a, out=out[0:2]) + + with pytest.raises(errors.DSLError, match="Subscripting a tuple"): + ProgramParser.apply_to_function(tuple_slice_program) + + +def test_vararg_per_element_domain_exception(vararg_scale_op): + scale = vararg_scale_op + + def vararg_domain_program( + tracers: tuple[gtx.Field[[IDim], float64], ...], + factor: float64, + out: tuple[gtx.Field[[IDim], float64], ...], + ): + scale(tracers, factor, out=out, domain=({IDim: (0, 2)}, {IDim: (0, 2)})) + + with pytest.raises(errors.DSLError) as exc_info: + ProgramParser.apply_to_function(vararg_domain_program) + + assert exc_info.match("Invalid call to 'scale'") + assert ( + re.search("Per-element domain tuples are not supported", exc_info.value.__cause__.args[0]) + is not None + ) + + def test_domain_exception_1(identity_def): domain_format_1 = gtx.field_operator(identity_def) From 264e753aaf6d8a268f10f9a2cc9dc9afaec9cfee Mon Sep 17 00:00:00 2001 From: Till Ehrengruber Date: Thu, 15 Oct 2026 00:35:19 +0200 Subject: [PATCH 16/24] refactor[next]: rename 'VarArgType' to 'VarLenTupleType' The type represents a homogeneous tuple whose length is not part of the type, not variadic function arguments (its historic use). Also: - 'VarArgElementKey' / 'VAR_ARG_ELEMENT' become 'VarLenTupleElementKey' / 'VAR_LENGTH_TUPLE_ELEMENT'. - '__str__' now prints 'tuple[T, ...]', matching both the annotation syntax users write and 'TupleType's 'tuple[a, b]'. - Test infrastructure: 'varargs_length' / 'DEFAULT_VARARGS_LENGTH' in 'cases.allocate' become 'tuple_length' / 'DEFAULT_TUPLE_LENGTH'; vararg-named unit tests aligned with the 'var_len_tuple' naming. --- src/gt4py/next/ffront/decorator.py | 2 +- .../ffront/foast_passes/type_deduction.py | 4 +- src/gt4py/next/ffront/foast_to_gtir.py | 4 +- src/gt4py/next/ffront/foast_to_past.py | 2 +- .../next/ffront/past_passes/type_deduction.py | 6 +-- src/gt4py/next/ffront/past_to_itir.py | 8 +-- src/gt4py/next/ffront/type_info.py | 4 +- src/gt4py/next/iterator/ir_utils/misc.py | 8 +-- src/gt4py/next/otf/compiled_program.py | 2 +- src/gt4py/next/type_system/type_info.py | 51 ++++++++++--------- .../next/type_system/type_specifications.py | 14 +++-- .../next/type_system/type_translation.py | 4 +- tests/next_tests/integration_tests/cases.py | 30 +++++------ .../feature_tests/ffront_tests/test_tuples.py | 14 ++--- .../ffront_tests/test_func_to_past.py | 26 +++++----- .../type_system_tests/test_type_info.py | 4 +- .../test_type_translation.py | 8 +-- 17 files changed, 99 insertions(+), 92 deletions(-) diff --git a/src/gt4py/next/ffront/decorator.py b/src/gt4py/next/ffront/decorator.py index 09b890413b..cd3a4b9f75 100644 --- a/src/gt4py/next/ffront/decorator.py +++ b/src/gt4py/next/ffront/decorator.py @@ -214,7 +214,7 @@ def _field_domain_descriptor_mapping_from_func_type(func_type: ts.FunctionType) for name, type_ in param_types.items(): for el_type_, path in type_info.primitive_constituents(type_, with_path_arg=True): if isinstance(el_type_, ts.FieldType): - if any(isinstance(idx, type_info.VarArgElementKey) for idx in path): + if any(isinstance(idx, type_info.VarLenTupleElementKey) for idx in path): raise NotImplementedError() path_as_expr = "".join(f"[{idx}]" for idx in path) static_domain_args.append(f"{name}{path_as_expr}") diff --git a/src/gt4py/next/ffront/foast_passes/type_deduction.py b/src/gt4py/next/ffront/foast_passes/type_deduction.py index bcea684b78..79bdcac56e 100644 --- a/src/gt4py/next/ffront/foast_passes/type_deduction.py +++ b/src/gt4py/next/ffront/foast_passes/type_deduction.py @@ -863,13 +863,13 @@ def visit_TupleComprehension( type=ts.TupleType(types=[new_mapper.element_expr.type for _ in element_types]), ) return result - elif isinstance(iterable.type, ts.VarArgType): + elif isinstance(iterable.type, ts.VarLenTupleType): 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_type = ts.VarLenTupleType(element_type=element_expr.type) return foast.TupleComprehension( inner=new_mapper, diff --git a/src/gt4py/next/ffront/foast_to_gtir.py b/src/gt4py/next/ffront/foast_to_gtir.py index 53b68a08cf..08ba691f7e 100644 --- a/src/gt4py/next/ffront/foast_to_gtir.py +++ b/src/gt4py/next/ffront/foast_to_gtir.py @@ -316,8 +316,8 @@ def visit_TupleComprehension(self, node: foast.TupleComprehension, **kwargs: Any 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) + assert isinstance(iterable_type, ts.VarLenTupleType) + assert isinstance(node.type, ts.VarLenTupleType) if isinstance(comprehension_target, itir.Sym): map_tuple_lambda = im.lambda_(comprehension_target)(element_expr) else: diff --git a/src/gt4py/next/ffront/foast_to_past.py b/src/gt4py/next/ffront/foast_to_past.py index fd51787e9b..f92e6221ff 100644 --- a/src/gt4py/next/ffront/foast_to_past.py +++ b/src/gt4py/next/ffront/foast_to_past.py @@ -117,7 +117,7 @@ def __call__(self, inp: ConcreteFOASTOperatorDef) -> ConcretePASTProgramDef: params_decl: list[past.Symbol] = [ past.DataSymbol( id=name, - type=type_, # note: this might be a generic type like VarArgType + type=type_, # note: this might be a generic type like VarLenTupleType namespace=dialect_ast_enums.Namespace.LOCAL, location=loc, ) diff --git a/src/gt4py/next/ffront/past_passes/type_deduction.py b/src/gt4py/next/ffront/past_passes/type_deduction.py index d282ccf7d7..a735833084 100644 --- a/src/gt4py/next/ffront/past_passes/type_deduction.py +++ b/src/gt4py/next/ffront/past_passes/type_deduction.py @@ -83,7 +83,7 @@ def _validate_field_dims(type_: ts.TypeSpec) -> None: ) elif isinstance(dom, past.TupleExpr): - if isinstance(out, ts.VarArgType): + if isinstance(out, ts.VarLenTupleType): # The fixed-length domain tuple would pin the length of the variable-length # output, so any call with a different length would break the program. raise ValueError( @@ -164,7 +164,7 @@ def visit_Slice(self, node: past.Slice, **kwargs: Any) -> past.Slice: def visit_Subscript(self, node: past.Subscript, **kwargs: Any) -> past.Subscript: value = self.visit(node.value, **kwargs) - if isinstance(value.type, (*ts.COLLECTION_TYPE_SPECS, ts.VarArgType)): + if isinstance(value.type, (*ts.COLLECTION_TYPE_SPECS, ts.VarLenTupleType)): # A subscript on a program argument is a domain slice; applying one to a # (possibly variable-length) tuple is not supported, and element selection is # not supported for any tuple argument. @@ -278,7 +278,7 @@ def visit_Call(self, node: past.Call, **kwargs: Any) -> past.Call: operator_return_type = type_info.return_type( new_func.type, with_args=arg_types, with_kwargs=kwarg_types ) - # Not equality: variadic comprehensions give a `VarArg[...]` return type, + # Not equality: variadic comprehensions give a variable-length tuple return type, # while 'out' is a concrete tuple. if not type_info.is_compatible_type(operator_return_type, new_kwargs["out"].type): raise ValueError( diff --git a/src/gt4py/next/ffront/past_to_itir.py b/src/gt4py/next/ffront/past_to_itir.py index aee05f0bdd..aed51892fb 100644 --- a/src/gt4py/next/ffront/past_to_itir.py +++ b/src/gt4py/next/ffront/past_to_itir.py @@ -211,14 +211,14 @@ def _compute_field_slice(node: past.Subscript) -> list[past.Slice]: def _unwrap_tuple_expr( - expr: past.Expr, path: tuple[int | type_info.VarArgElementKey, ...] -) -> tuple[past.Expr, Sequence[int | type_info.VarArgElementKey]]: + expr: past.Expr, path: tuple[int | type_info.VarLenTupleElementKey, ...] +) -> tuple[past.Expr, Sequence[int | type_info.VarLenTupleElementKey]]: """Unwrap (nested) TupleExpr by following the given path as long as possible. If a non-tuple expression is encountered, the current expression and the remaining path are returned. """ - path_remainder: Sequence[int | type_info.VarArgElementKey] = path + path_remainder: Sequence[int | type_info.VarLenTupleElementKey] = path while isinstance(expr, past.TupleExpr): idx, *path_remainder = path_remainder # a `TupleExpr` always has a fixed-length tuple type, so the path element is an index @@ -472,7 +472,7 @@ def _visit_stencil_call_out_arg( def impl( out_type: ts.FieldType, current: itir.Expr, - path: tuple[int | type_info.VarArgElementKey, ...], + path: tuple[int | type_info.VarLenTupleElementKey, ...], ) -> itir.Expr: out_field, path_remainder = _unwrap_tuple_expr(out_arg, path) diff --git a/src/gt4py/next/ffront/type_info.py b/src/gt4py/next/ffront/type_info.py index 8eb3a516ae..ffdf7b63bf 100644 --- a/src/gt4py/next/ffront/type_info.py +++ b/src/gt4py/next/ffront/type_info.py @@ -153,9 +153,9 @@ def function_signature_incompatibilities_fieldop( def _tree_map_type_constructor_drop_python_type( - value: ts.CollectionTypeSpecT | ts.VarArgType, + value: ts.CollectionTypeSpecT | ts.VarLenTupleType, elems: Iterable[ts.DataType | ts.DimensionType | ts.DeferredType], -) -> ts.CollectionTypeSpecT | ts.VarArgType: +) -> ts.CollectionTypeSpecT | ts.VarLenTupleType: result = type_info.tree_map_type_constructor(value, elems) if isinstance(value, ts.NamedCollectionType): result = datamodels.evolve(result, original_python_type=ts.ANY_PYTHON_TYPE_NAME) diff --git a/src/gt4py/next/iterator/ir_utils/misc.py b/src/gt4py/next/iterator/ir_utils/misc.py index 3ae0af8bf8..7b8a9d75fd 100644 --- a/src/gt4py/next/iterator/ir_utils/misc.py +++ b/src/gt4py/next/iterator/ir_utils/misc.py @@ -300,18 +300,18 @@ def tree_map_type_into_expr( Like `type_info.tree_map_type`, but `fun` additionally receives a `current` argument — a reference to the value of the current element — and the mapped callable is called with the value expression alongside the type. Fixed-length collection nodes recombine - the element results with `make_tuple`, rebinding `current` per element; `VarArgType` + the element results with `make_tuple`, rebinding `current` per element; `VarLenTupleType` nodes map their elements with the `map_tuple` builtin, so the result is also valid for types whose concrete length is only known later. The bindings introduced along the way are inlined unconditionally before returning — the value expression should hence be cheap (e.g. a symbol reference or tuple accesses thereon), as it may be duplicated. - With `with_path_arg`, paths contain `type_info.VAR_ARG_ELEMENT` below variable-length + With `with_path_arg`, paths contain `type_info.VAR_LENGTH_TUPLE_ELEMENT` below variable-length tuples. Can be used directly or as a (possibly parametrized) decorator. >>> field = ts.FieldType(dims=[], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64)) - >>> type_ = ts.TupleType(types=[field, ts.VarArgType(element_type=field)]) + >>> type_ = ts.TupleType(types=[field, ts.VarLenTupleType(element_type=field)]) >>> expr = tree_map_type_into_expr(lambda t, current: im.call("get_dom")(current))( ... type_, im.ref("out") ... ) @@ -328,7 +328,7 @@ def impl(type_: ts.TypeSpec, expr: itir.Expr) -> itir.Expr | tuple[itir.Expr, .. current = im.ref(_CURRENT_ELEMENT) def collection_constructor(value: ts.TypeSpec, elts: Iterable[itir.Expr]) -> itir.Expr: - if isinstance(value, ts.VarArgType): + if isinstance(value, ts.VarLenTupleType): (element_expr,) = elts return im.call(im.call("map_tuple")(im.lambda_(_CURRENT_ELEMENT)(element_expr)))( current diff --git a/src/gt4py/next/otf/compiled_program.py b/src/gt4py/next/otf/compiled_program.py index 99d9ac1edf..c9700ce8b1 100644 --- a/src/gt4py/next/otf/compiled_program.py +++ b/src/gt4py/next/otf/compiled_program.py @@ -515,7 +515,7 @@ def _is_generic(self) -> bool: # in the type system. For now we just use `DeferredType` to communicate between # here and `type_info.type_in_program_context`. def is_generic_type(t: ts.TypeSpec) -> bool: - if isinstance(t, (ts.DeferredType, ts.VarArgType)): + if isinstance(t, (ts.DeferredType, ts.VarLenTupleType)): return True if isinstance(t, ts.TupleType): return any(is_generic_type(el) for el in t.types) diff --git a/src/gt4py/next/type_system/type_info.py b/src/gt4py/next/type_system/type_info.py index 023bb080f8..6a2080690c 100644 --- a/src/gt4py/next/type_system/type_info.py +++ b/src/gt4py/next/type_system/type_info.py @@ -84,17 +84,17 @@ def type_class(symbol_type: ts.TypeSpec) -> Type[ts.TypeSpec]: ) -class VarArgElementKey: - """Path key of the representative element of a `VarArgType` in type traversals.""" +class VarLenTupleElementKey: + """Path key of the representative element of a `VarLenTupleType` in type traversals.""" def __repr__(self) -> str: - return "VAR_ARG_ELEMENT" + return "VAR_LENGTH_TUPLE_ELEMENT" -#: Type traversals (`tree_map_type`, `primitive_constituents`) visit a `VarArgType` as +#: Type traversals (`tree_map_type`, `primitive_constituents`) visit a `VarLenTupleType` as #: a single representative element (its `element_type`), addressed by this sentinel #: instead of a position. -VAR_ARG_ELEMENT: Final = VarArgElementKey() +VAR_LENGTH_TUPLE_ELEMENT: Final = VarLenTupleElementKey() @overload @@ -106,20 +106,20 @@ def primitive_constituents( @overload def primitive_constituents( symbol_type: ts.TypeSpec, with_path_arg: Literal[True] -) -> utils.XIterable[tuple[ts.TypeSpec, tuple[int | VarArgElementKey, ...]]]: ... +) -> utils.XIterable[tuple[ts.TypeSpec, tuple[int | VarLenTupleElementKey, ...]]]: ... def primitive_constituents( symbol_type: ts.TypeSpec, with_path_arg: bool = False ) -> ( utils.XIterable[ts.TypeSpec] - | utils.XIterable[tuple[ts.TypeSpec, tuple[int | VarArgElementKey, ...]]] + | utils.XIterable[tuple[ts.TypeSpec, tuple[int | VarLenTupleElementKey, ...]]] ): """ Return the primitive types contained in a composite type. - A `VarArgType` is traversed as its single representative element, addressed by - `VAR_ARG_ELEMENT` in paths. + A `VarLenTupleType` is traversed as its single representative element, addressed by + `VAR_LENGTH_TUPLE_ELEMENT` in paths. >>> from gt4py.next import common >>> I = common.Dimension(value="I") @@ -136,10 +136,15 @@ def primitive_constituents( """ def constituents_yielder( - symbol_type: ts.TypeSpec, path: tuple[int | VarArgElementKey, ...] - ) -> Iterator[ts.TypeSpec] | Iterator[tuple[ts.TypeSpec, tuple[int | VarArgElementKey, ...]]]: - if isinstance(symbol_type, ts.VarArgType): - yield from constituents_yielder(symbol_type.element_type, (*path, VAR_ARG_ELEMENT)) + symbol_type: ts.TypeSpec, path: tuple[int | VarLenTupleElementKey, ...] + ) -> ( + Iterator[ts.TypeSpec] + | Iterator[tuple[ts.TypeSpec, tuple[int | VarLenTupleElementKey, ...]]] + ): + if isinstance(symbol_type, ts.VarLenTupleType): + yield from constituents_yielder( + symbol_type.element_type, (*path, VAR_LENGTH_TUPLE_ELEMENT) + ) elif isinstance(symbol_type, ts.COLLECTION_TYPE_SPECS): for i, el_type in enumerate(symbol_type.types): yield from constituents_yielder(el_type, (*path, i)) @@ -157,26 +162,26 @@ def constituents_yielder( def _tree_map_type_collection_elements(value: ts.TypeSpec) -> Iterable[ts.TypeSpec]: - if isinstance(value, ts.VarArgType): + if isinstance(value, ts.VarLenTupleType): return (value.element_type,) assert isinstance(value, ts.COLLECTION_TYPE_SPECS) return value.types def _tree_map_type_collection_keys(value: ts.TypeSpec) -> Iterable[Any]: - if isinstance(value, ts.VarArgType): - return (VAR_ARG_ELEMENT,) + if isinstance(value, ts.VarLenTupleType): + return (VAR_LENGTH_TUPLE_ELEMENT,) return itertools.count() def tree_map_type_constructor( - value: ts.CollectionTypeSpecT | ts.VarArgType, + value: ts.CollectionTypeSpecT | ts.VarLenTupleType, elems: Iterable[ts.DataType | ts.DimensionType | ts.DeferredType], -) -> ts.CollectionTypeSpecT | ts.VarArgType: - if isinstance(value, ts.VarArgType): +) -> ts.CollectionTypeSpecT | ts.VarLenTupleType: + if isinstance(value, ts.VarLenTupleType): (element_type,) = elems assert isinstance(element_type, ts.DataType) - return ts.VarArgType(element_type=element_type) + return ts.VarLenTupleType(element_type=element_type) return ( ts.NamedCollectionType( keys=value.keys, original_python_type=value.original_python_type, types=list(elems) @@ -233,7 +238,7 @@ def tree_map_type( ) return next_utils.tree_map( fun, - collection_type=(*ts.COLLECTION_TYPE_SPECS, ts.VarArgType), + collection_type=(*ts.COLLECTION_TYPE_SPECS, ts.VarLenTupleType), collection_elements=_tree_map_type_collection_elements, collection_keys=_tree_map_type_collection_keys, result_collection_constructor=result_collection_constructor, @@ -614,9 +619,9 @@ def is_concretizable(symbol_type: ts.TypeSpec, to_type: ts.TypeSpec) -> bool: or issubclass(type_class(to_type), symbol_type.constraint) ): return True - if isinstance(symbol_type, ts.VarArgType) and isinstance(to_type, ts.VarArgType): + if isinstance(symbol_type, ts.VarLenTupleType) and isinstance(to_type, ts.VarLenTupleType): return is_concretizable(symbol_type.element_type, to_type.element_type) - if isinstance(symbol_type, ts.VarArgType) and isinstance(to_type, ts.TupleType): + if isinstance(symbol_type, ts.VarLenTupleType) and isinstance(to_type, ts.TupleType): # The empty tuple is a valid concretization of any variable-length tuple: there # are no elements whose type could conflict with the element type. return len(to_type.types) == 0 or ( diff --git a/src/gt4py/next/type_system/type_specifications.py b/src/gt4py/next/type_system/type_specifications.py index 73a8a08a5d..91edc45457 100644 --- a/src/gt4py/next/type_system/type_specifications.py +++ b/src/gt4py/next/type_system/type_specifications.py @@ -142,14 +142,18 @@ def __str__(self) -> str: return f"tuple[{', '.join(map(str, self.types))}]" -# TODO(tehrengruber): Rename to var length tuple or similar -class VarArgType(DataType): - """Represents a variable number of arguments of the same type.""" +class VarLenTupleType(DataType): + """ + A homogeneous tuple whose length is not part of the type. + + The length is only known at the call, where the type is specialized to a concrete + `TupleType` (or stays variable-length in length-agnostic code paths). + """ element_type: DataType def __str__(self) -> str: - return f"VarArg[{self.element_type}]" + return f"tuple[{self.element_type}, ...]" class AnyPythonType: @@ -195,7 +199,7 @@ def __str__(self) -> str: return f"NamedTuple{{{', '.join(f'{k}: {v}' for k, v in zip(self.keys, self.types))}}}" -# TODO(tehrengruber): Add VarArgType to this list, which requires checking all users of it +# TODO(tehrengruber): Add VarLenTupleType to this list, which requires checking all users of it # that potentially depend on the length being known. CollectionTypeSpecT = TypeVar("CollectionTypeSpecT", TupleType, NamedCollectionType) CollectionTypeSpec = TupleType | NamedCollectionType diff --git a/src/gt4py/next/type_system/type_translation.py b/src/gt4py/next/type_system/type_translation.py index 698bb8c7c1..aa44b1ee99 100644 --- a/src/gt4py/next/type_system/type_translation.py +++ b/src/gt4py/next/type_system/type_translation.py @@ -81,7 +81,7 @@ def make_constructor_type(type_spec: ts.TypeSpec) -> ts.ConstructorType: pos_only_args=[ts.DeferredType(constraint=None)], pos_or_kw_args={}, kw_only_args={}, - returns=ts.DeferredType(constraint=ts.VarArgType), + returns=ts.DeferredType(constraint=ts.VarLenTupleType), ) ) @@ -235,7 +235,7 @@ def from_type_hint( raise ValueError( f"Tuple annotation '{type_hint}' contains an invalid element type." ) - return ts.VarArgType(element_type=element_type) + return ts.VarLenTupleType(element_type=element_type) # Unparametrized annotation, i.e. bare `tuple` / `typing.Tuple` (the # empty-tuple annotation `tuple[()]` has `args == ()` and is rejected below) elif args is None: diff --git a/tests/next_tests/integration_tests/cases.py b/tests/next_tests/integration_tests/cases.py index 8577243008..a1c39e6f90 100644 --- a/tests/next_tests/integration_tests/cases.py +++ b/tests/next_tests/integration_tests/cases.py @@ -347,7 +347,7 @@ def allocate( strategy: Optional[DataInitializer] = None, dtype: Optional[np.typing.DTypeLike] = None, extend: Optional[dict[gtx.Dimension, tuple[int, int]]] = None, - varargs_length: Optional[int] = None, + tuple_length: Optional[int] = None, ) -> FieldViewArg: """ Allocate a parameter or return value from a fieldview code object. @@ -365,8 +365,8 @@ def allocate( extend: Lower and upper size extension per dimension. Useful for shifted fields, which must start off bigger than the output field in the shifted dimension. - varargs_length: Concrete length to allocate for variable-length tuple - parameters (default: ``DEFAULT_VARARGS_LENGTH``). + tuple_length: Concrete length to allocate for variable-length tuple + parameters (default: ``DEFAULT_TUPLE_LENGTH``). """ if sizes: assert not domain and all(dim in case.default_sizes for dim in sizes) @@ -396,7 +396,7 @@ def allocate( domain=domain, dtype=dtype, strategy=strategy.from_case(case=case, fieldview_prog=fieldview_prog, arg_name=name), - varargs_length=varargs_length, + tuple_length=tuple_length, ) @@ -594,7 +594,7 @@ def unstructured_case_3d(unstructured_case): #: Default concrete length used for variable-length tuple parameters, e.g. when #: allocating arguments for them. -DEFAULT_VARARGS_LENGTH: Final = 3 +DEFAULT_TUPLE_LENGTH: Final = 3 def _allocate_from_type( @@ -604,7 +604,7 @@ def _allocate_from_type( strategy: DataInitializer, dtype: Optional[np.typing.DTypeLike] = None, tuple_start: Optional[int] = None, - varargs_length: Optional[int] = None, + tuple_length: Optional[int] = None, ) -> FieldViewArg: """Allocate data based on the type or a (nested) tuple thereof.""" match arg_type: @@ -625,16 +625,16 @@ def _allocate_from_type( domain=domain, dtype=dtype, strategy=strategy, - varargs_length=varargs_length, + tuple_length=tuple_length, ) for t in types ) ) - case ts.VarArgType(element_type=element_type): - length = varargs_length if varargs_length is not None else DEFAULT_VARARGS_LENGTH + case ts.VarLenTupleType(element_type=element_type): + length = tuple_length if tuple_length is not None else DEFAULT_TUPLE_LENGTH # the unique-value offsets are computed from the declared types using the # default length (see `get_param_size`) - assert length <= DEFAULT_VARARGS_LENGTH + assert length <= DEFAULT_TUPLE_LENGTH return tuple( ( _allocate_from_type( @@ -643,7 +643,7 @@ def _allocate_from_type( domain=domain, dtype=dtype, strategy=strategy, - varargs_length=varargs_length, + tuple_length=tuple_length, ) for t in [element_type] * length ) @@ -663,7 +663,7 @@ def _allocate_from_type( domain=domain, dtype=dtype, strategy=strategy, - varargs_length=varargs_length, + tuple_length=tuple_length, ) for t in types ) @@ -698,10 +698,8 @@ def get_param_size(param_type: ts.TypeSpec, sizes: dict[gtx.Dimension, int]) -> return sum([get_param_size(t, sizes=sizes) for t in types]) case ts.NamedCollectionType(types=types): return sum([get_param_size(t, sizes=sizes) for t in types]) - case ts.VarArgType(element_type=element_type): - return get_param_size( - ts.TupleType(types=[element_type] * DEFAULT_VARARGS_LENGTH), sizes - ) + case ts.VarLenTupleType(element_type=element_type): + return get_param_size(ts.TupleType(types=[element_type] * DEFAULT_TUPLE_LENGTH), sizes) case _: raise TypeError(f"Can not get size for parameter of type '{param_type}'.") diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py index 63983ead54..487558cb91 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_tuples.py @@ -189,8 +189,8 @@ def testee(tracers: tuple[cases.IField, ...], factor: int32) -> tuple[cases.IFie return tuple(tracer * factor for tracer in tracers) for length in (3, 2): - tracers = cases.allocate(cartesian_case, testee, "tracers", varargs_length=length)() - out = cases.allocate(cartesian_case, testee, cases.RETURN, varargs_length=length)() + tracers = cases.allocate(cartesian_case, testee, "tracers", tuple_length=length)() + out = cases.allocate(cartesian_case, testee, cases.RETURN, tuple_length=length)() cases.verify( cartesian_case, testee, @@ -212,8 +212,8 @@ def testee(tracers: tuple[cases.IField, ...], factor: int32, out: tuple[cases.IF scale(tracers, factor, out=out) for length in (3, 2): # ensure each concrete length gets its own specialization - tracers = cases.allocate(cartesian_case, testee, "tracers", varargs_length=length)() - out = cases.allocate(cartesian_case, testee, "out", varargs_length=length)() + tracers = cases.allocate(cartesian_case, testee, "tracers", tuple_length=length)() + out = cases.allocate(cartesian_case, testee, "out", tuple_length=length)() cases.run(cartesian_case, testee, tracers, 42, out) assert all(np.allclose(o.asnumpy(), t.asnumpy() * 42) for o, t in zip(out, tracers)) @@ -360,9 +360,9 @@ def testee( ) -> tuple[cases.IField, ...]: return a if c else b - a = cases.allocate(cartesian_case, testee, "a", varargs_length=2)() - b = cases.allocate(cartesian_case, testee, "b", varargs_length=3)() - out = cases.allocate(cartesian_case, testee, cases.RETURN, varargs_length=2)() + a = cases.allocate(cartesian_case, testee, "a", tuple_length=2)() + b = cases.allocate(cartesian_case, testee, "b", tuple_length=3)() + out = cases.allocate(cartesian_case, testee, cases.RETURN, tuple_length=2)() with pytest.raises(errors.DSLError): cases.run(cartesian_case, testee, a, b, False, out=out) diff --git a/tests/next_tests/unit_tests/ffront_tests/test_func_to_past.py b/tests/next_tests/unit_tests/ffront_tests/test_func_to_past.py index 7cc86a5a0b..533113bfc2 100644 --- a/tests/next_tests/unit_tests/ffront_tests/test_func_to_past.py +++ b/tests/next_tests/unit_tests/ffront_tests/test_func_to_past.py @@ -153,7 +153,7 @@ def test_copy_restrict_parsing(copy_restrict_program_def): @pytest.fixture -def vararg_scale_op(): +def var_len_tuple_scale_op(): @gtx.field_operator def scale( tracers: tuple[gtx.Field[[IDim], float64], ...], factor: float64 @@ -163,10 +163,10 @@ def scale( return scale -def test_vararg_subscript_out_exception(vararg_scale_op): - scale = vararg_scale_op +def test_var_len_tuple_subscript_out_exception(var_len_tuple_scale_op): + scale = var_len_tuple_scale_op - def vararg_slice_program( + def var_len_tuple_slice_program( tracers: tuple[gtx.Field[[IDim], float64], ...], factor: float64, out: tuple[gtx.Field[[IDim], float64], ...], @@ -174,13 +174,13 @@ def vararg_slice_program( scale(tracers, factor, out=out[0:2]) with pytest.raises(errors.DSLError, match="Subscripting a tuple"): - ProgramParser.apply_to_function(vararg_slice_program) + ProgramParser.apply_to_function(var_len_tuple_slice_program) -def test_vararg_subscript_arg_exception(vararg_scale_op): - scale = vararg_scale_op +def test_var_len_tuple_subscript_arg_exception(var_len_tuple_scale_op): + scale = var_len_tuple_scale_op - def vararg_slice_program( + def var_len_tuple_slice_program( tracers: tuple[gtx.Field[[IDim], float64], ...], factor: float64, out: tuple[gtx.Field[[IDim], float64], ...], @@ -188,7 +188,7 @@ def vararg_slice_program( scale(tracers[0:2], factor, out=out) with pytest.raises(errors.DSLError, match="Subscripting a tuple"): - ProgramParser.apply_to_function(vararg_slice_program) + ProgramParser.apply_to_function(var_len_tuple_slice_program) def test_tuple_subscript_out_exception(): @@ -208,10 +208,10 @@ def tuple_slice_program( ProgramParser.apply_to_function(tuple_slice_program) -def test_vararg_per_element_domain_exception(vararg_scale_op): - scale = vararg_scale_op +def test_var_len_tuple_per_element_domain_exception(var_len_tuple_scale_op): + scale = var_len_tuple_scale_op - def vararg_domain_program( + def var_len_tuple_domain_program( tracers: tuple[gtx.Field[[IDim], float64], ...], factor: float64, out: tuple[gtx.Field[[IDim], float64], ...], @@ -219,7 +219,7 @@ def vararg_domain_program( scale(tracers, factor, out=out, domain=({IDim: (0, 2)}, {IDim: (0, 2)})) with pytest.raises(errors.DSLError) as exc_info: - ProgramParser.apply_to_function(vararg_domain_program) + ProgramParser.apply_to_function(var_len_tuple_domain_program) assert exc_info.match("Invalid call to 'scale'") assert ( diff --git a/tests/next_tests/unit_tests/type_system_tests/test_type_info.py b/tests/next_tests/unit_tests/type_system_tests/test_type_info.py index 019d4cf648..1d52ba2e38 100644 --- a/tests/next_tests/unit_tests/type_system_tests/test_type_info.py +++ b/tests/next_tests/unit_tests/type_system_tests/test_type_info.py @@ -467,7 +467,7 @@ def test_needs_value_extraction(type_spec: ts.TypeSpec, expected: bool): def test_is_concretizable_vararg_to_tuple(): float_type = ts.ScalarType(kind=ts.ScalarKind.FLOAT64) int_type = ts.ScalarType(kind=ts.ScalarKind.INT32) - vararg = ts.VarArgType(element_type=float_type) + vararg = ts.VarLenTupleType(element_type=float_type) assert type_info.is_concretizable(vararg, ts.TupleType(types=[float_type, float_type])) assert not type_info.is_concretizable(vararg, ts.TupleType(types=[int_type])) @@ -496,7 +496,7 @@ def test_promote_lists(): def test_is_concretizable_tuple_with_nested_vararg(): float_type = ts.ScalarType(kind=ts.ScalarKind.FLOAT64) field_type = ts.FieldType(dims=[], dtype=float_type) - nested = ts.TupleType(types=[ts.VarArgType(element_type=field_type), field_type]) + nested = ts.TupleType(types=[ts.VarLenTupleType(element_type=field_type), field_type]) assert type_info.is_concretizable( nested, ts.TupleType(types=[ts.TupleType(types=[field_type, field_type]), field_type]) diff --git a/tests/next_tests/unit_tests/type_system_tests/test_type_translation.py b/tests/next_tests/unit_tests/type_system_tests/test_type_translation.py index 8778e0a18e..6d73512150 100644 --- a/tests/next_tests/unit_tests/type_system_tests/test_type_translation.py +++ b/tests/next_tests/unit_tests/type_system_tests/test_type_translation.py @@ -218,7 +218,7 @@ def test_generic_tuple_constructor_type(): pos_only_args=[ts.DeferredType(constraint=None)], pos_or_kw_args={}, kw_only_args={}, - returns=ts.DeferredType(constraint=ts.VarArgType), + returns=ts.DeferredType(constraint=ts.VarLenTupleType), ) ) assert ( @@ -249,11 +249,11 @@ def test_invalid_symbol_types(): with pytest.raises(ValueError, match="invalid element type"): type_translation.from_type_hint(tuple[tuple, ...]) - # Variadic tuples (`tuple[T, ...]`) are now valid — returns a VarArgType. - assert type_translation.from_type_hint(tuple[int, ...]) == ts.VarArgType( + # Variadic tuples (`tuple[T, ...]`) are now valid — returns a VarLenTupleType. + assert type_translation.from_type_hint(tuple[int, ...]) == ts.VarLenTupleType( element_type=ts.ScalarType(kind=ts.ScalarKind.INT64) ) - assert type_translation.from_type_hint(typing.Tuple["float", ...]) == ts.VarArgType( + assert type_translation.from_type_hint(typing.Tuple["float", ...]) == ts.VarLenTupleType( element_type=ts.ScalarType(kind=ts.ScalarKind.FLOAT64) ) From 8bb281388a5c2e777dd2889ffa615e9d5e21132c Mon Sep 17 00:00:00 2001 From: Till Ehrengruber Date: Thu, 15 Oct 2026 01:12:27 +0200 Subject: [PATCH 17/24] fix[next]: don't assert argument compatibility for operator wrappers For the generated wrapper programs of direct operator calls the parameter specialization in 'OperatorToProgram' runs before any argument validation, so incompatible arguments (e.g. a 'float32' field passed to an 'int32' operator) tripped the assert in '_specialize_program_params' instead of surfacing as a 'DSLError'. Disable the compatibility check on that path; the type deduction of the wrapper body call right after rejects incompatible arguments with the proper user-facing error. --- src/gt4py/next/ffront/foast_to_past.py | 7 ++++++- src/gt4py/next/ffront/past_to_itir.py | 9 +++++++-- 2 files changed, 13 insertions(+), 3 deletions(-) diff --git a/src/gt4py/next/ffront/foast_to_past.py b/src/gt4py/next/ffront/foast_to_past.py index f92e6221ff..28d8c5a042 100644 --- a/src/gt4py/next/ffront/foast_to_past.py +++ b/src/gt4py/next/ffront/foast_to_past.py @@ -164,7 +164,12 @@ def __call__(self, inp: ConcreteFOASTOperatorDef) -> ConcretePASTProgramDef: # this step would become obsolete, as the generic variant could be type checked # directly. untyped_past_node = past_to_itir._specialize_program_params( - untyped_past_node, inp.args.args, inp.args.kwargs + untyped_past_node, + inp.args.args, + inp.args.kwargs, + # argument validation has not run yet on this path; incompatible arguments are + # rejected right below, by the type deduction of the body call + check_compatibility=False, ) untyped_past_node = closure_var_type_deduction.ClosureVarTypeDeduction.apply( untyped_past_node, diff --git a/src/gt4py/next/ffront/past_to_itir.py b/src/gt4py/next/ffront/past_to_itir.py index aed51892fb..911af82a43 100644 --- a/src/gt4py/next/ffront/past_to_itir.py +++ b/src/gt4py/next/ffront/past_to_itir.py @@ -229,7 +229,10 @@ def _unwrap_tuple_expr( def _specialize_program_params( - past_node: past.Program, arg_types: Sequence[ts.TypeSpec], kwarg_types: dict[str, ts.TypeSpec] + past_node: past.Program, + arg_types: Sequence[ts.TypeSpec], + kwarg_types: dict[str, ts.TypeSpec], + check_compatibility: bool = True, ) -> past.Program: """ Replace the program parameter types by the concrete argument types of the call. @@ -246,7 +249,9 @@ def _specialize_program_params( for param in past_node.params: param_type = param_types[str(param.id)] # incompatible argument types have already been rejected in `transform_program_args` - assert type_info.is_concretizable(param.type, to_type=param_type), ( + assert not check_compatibility or type_info.is_concretizable( + param.type, to_type=param_type + ), ( f"Argument of type '{param_type}' is not compatible with parameter " f"'{param.id}' of type '{param.type}'." ) From 66dabcecb628e92022097679af9459407009df6b Mon Sep 17 00:00:00 2001 From: Till Ehrengruber Date: Thu, 15 Oct 2026 10:30:31 +0200 Subject: [PATCH 18/24] fix[next]: lower uncalled programs with their declared parameter types MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The 'Program.gtir' property lowers a program without a call (used e.g. by the input / output name extractors and in tests); it passed empty 'CompileTimeArgs' which broke the parameter specialization in 'past_to_gtir'. Use the declared parameter types as compile-time arguments instead — the specialization is then the identity. --- src/gt4py/next/ffront/decorator.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/src/gt4py/next/ffront/decorator.py b/src/gt4py/next/ffront/decorator.py index cd3a4b9f75..c94341b7da 100644 --- a/src/gt4py/next/ffront/decorator.py +++ b/src/gt4py/next/ffront/decorator.py @@ -308,15 +308,20 @@ def _all_closure_vars(self) -> dict[str, Any]: @functools.cached_property def gtir(self) -> itir.Program: - no_args_past = workflow.ConcreteArtifact( + # The program is lowered without a call here, so use the declared parameter + # types as compile-time arguments. + declared_args_past = workflow.ConcreteArtifact( data=ffront_stages.PASTProgramDef( past_node=self.past_stage.past_node, closure_vars=self.past_stage.closure_vars, grid_type=self.definition_stage.grid_type, ), - args=arguments.CompileTimeArgs.empty(), + args=dataclasses.replace( + arguments.CompileTimeArgs.empty(), + args=tuple(param.type for param in self.past_stage.past_node.params), + ), ) - return self._frontend_transforms.past_to_itir(no_args_past).data + return self._frontend_transforms.past_to_itir(declared_args_past).data def with_grid_type(self, grid_type: common.GridType) -> Program: return dataclasses.replace( From 3c61f57159a3d491d14b3e2bc7d9d17911f4fedb Mon Sep 17 00:00:00 2001 From: Till Ehrengruber Date: Fri, 16 Oct 2026 13:38:32 +0200 Subject: [PATCH 19/24] Cleanup comments --- src/gt4py/next/ffront/foast_to_past.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/src/gt4py/next/ffront/foast_to_past.py b/src/gt4py/next/ffront/foast_to_past.py index 28d8c5a042..fa44a99bbd 100644 --- a/src/gt4py/next/ffront/foast_to_past.py +++ b/src/gt4py/next/ffront/foast_to_past.py @@ -102,15 +102,15 @@ def __call__(self, inp: ConcreteFOASTOperatorDef) -> ConcretePASTProgramDef: # with the out argument of the program we generate here. loc = inp.data.foast_node.location assert isinstance(inp.data.foast_node.type, ts.CallableType) - partial_program_type = ffront_type_info.type_in_program_context(inp.data.foast_node.type) - assert isinstance(partial_program_type, ts_ffront.ProgramType) + program_type = ffront_type_info.type_in_program_context(inp.data.foast_node.type) + assert isinstance(program_type, ts_ffront.ProgramType) param_types: tuple[tuple[str, ts.TypeSpec], ...] = ( *( (f"__pos_arg{i}", type_) - for i, type_ in enumerate(partial_program_type.definition.pos_only_args) + for i, type_ in enumerate(program_type.definition.pos_only_args) ), - *partial_program_type.definition.pos_or_kw_args.items(), - *partial_program_type.definition.kw_only_args.items(), + *program_type.definition.pos_or_kw_args.items(), + *program_type.definition.kw_only_args.items(), ) assert param_types[-1][0] == "out" @@ -163,12 +163,15 @@ def __call__(self, inp: ConcreteFOASTOperatorDef) -> ConcretePASTProgramDef: # specialized here. With support for generic types in the program type deduction # this step would become obsolete, as the generic variant could be type checked # directly. + # TODO(tehrengruber): Remove again. This specialization is only needed as the `program_type` + # above contains `DeferredType` for direct scan operator calls due to the lack of + # generics. untyped_past_node = past_to_itir._specialize_program_params( untyped_past_node, inp.args.args, inp.args.kwargs, # argument validation has not run yet on this path; incompatible arguments are - # rejected right below, by the type deduction of the body call + # rejected right below with the proper error message check_compatibility=False, ) untyped_past_node = closure_var_type_deduction.ClosureVarTypeDeduction.apply( From 0d80c32ff641339bc50cb9306e1116e24d2e72fa Mon Sep 17 00:00:00 2001 From: Till Ehrengruber Date: Fri, 16 Oct 2026 14:36:11 +0200 Subject: [PATCH 20/24] fix[next]: derive named-collection member types from runtime values 'from_value' typed dataclass / NamedTuple instances from the class annotations, so a member declared as a variable-length tuple (e.g. 'tracers: tuple[IField, ...]') stayed generic in the call-time argument types and the lowering failed. Take the member types from the runtime member values instead (consistent with plain tuples); in debug mode additionally check them for compatibility with the member annotations. Note that members holding plain Python scalar literals are now typed by the 'from_value' convention (e.g. 'float64' for floats) instead of their annotation, consistent with regular parameters, where literals of a non-default dtype are already rejected (pre-existing limitation). --- .../next/type_system/type_translation.py | 26 ++++++++- .../ffront_tests/test_named_collections.py | 43 ++++++++++++++- .../test_type_translation.py | 54 +++++++++++++++++++ 3 files changed, 120 insertions(+), 3 deletions(-) diff --git a/src/gt4py/next/type_system/type_translation.py b/src/gt4py/next/type_system/type_translation.py index aa44b1ee99..afeaf254e2 100644 --- a/src/gt4py/next/type_system/type_translation.py +++ b/src/gt4py/next/type_system/type_translation.py @@ -19,7 +19,7 @@ import sys import types import typing -from typing import Any, ForwardRef, Optional, TypeAlias +from typing import Any, ForwardRef, Optional, TypeAlias, cast import numpy as np import numpy.typing as npt @@ -405,6 +405,30 @@ def from_value(value: Any) -> ts.TypeSpec: elems = [from_value(el) for el in value] assert all(isinstance(elem, ts.DataType) for elem in elems) return ts.TupleType(types=elems) + elif isinstance(value, named_collections.CUSTOM_NAMED_COLLECTION_TYPES): + # Take the member types from the runtime values — the annotated types may be + # generic, e.g. a variable-length tuple, whose concrete type is only known from + # the value. Members whose type is incompatible with the annotation are rejected + # later, by the signature checks against the (annotated) parameter types; in + # debug mode reject them already here, with a more precise error. + annotated_type = from_type_hint(xtyping.infer_type(value)) + assert isinstance(annotated_type, ts.NamedCollectionType) + member_types = [from_value(getattr(value, key)) for key in annotated_type.keys] + if __debug__: + for key, annotated_member_type, member_type in zip( + annotated_type.keys, annotated_type.types, member_types, strict=True + ): + if not type_info.is_concretizable(annotated_member_type, to_type=member_type): + raise ValueError( + f"Member '{key}' of '{type(value).__name__}' has type '{member_type}'," + f" which is incompatible with its annotation '{annotated_member_type}'." + ) + assert all(isinstance(t, (ts.DataType, ts.DimensionType)) for t in member_types) + return ts.NamedCollectionType( + types=cast(list[ts.DataType], member_types), + keys=list(annotated_type.keys), + original_python_type=annotated_type.original_python_type, + ) elif isinstance(value, PythonNamespaceObject): return NamespaceProxy(value) else: diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_named_collections.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_named_collections.py index 911ace0804..302f03c63b 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_named_collections.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_named_collections.py @@ -110,6 +110,45 @@ def test_named_tuple_like_constructed_outside(cartesian_case, testee): ) +class VarLenMemberNamedTuple(NamedTuple): + tracers: tuple[gtx.Field[[IDim, JDim], gtx.float32], ...] + rho: gtx.Field[[IDim, JDim], gtx.float32] + + +@dataclasses.dataclass +class VarLenMemberDataclass: + tracers: tuple[gtx.Field[[IDim, JDim], gtx.float32], ...] + rho: gtx.Field[[IDim, JDim], gtx.float32] + + +@gtx.field_operator +def add_rho_named_tuple( + c: VarLenMemberNamedTuple, +) -> tuple[gtx.Field[[IDim, JDim], gtx.float32], ...]: + return tuple(tracer + c.rho for tracer in c.tracers) + + +@gtx.field_operator +def add_rho_dataclass( + c: VarLenMemberDataclass, +) -> tuple[gtx.Field[[IDim, JDim], gtx.float32], ...]: + return tuple(tracer + c.rho for tracer in c.tracers) + + +@pytest.mark.parametrize("testee", [add_rho_named_tuple, add_rho_dataclass]) +@pytest.mark.uses_tuple_args +def test_named_collection_with_var_len_tuple_member(cartesian_case, testee): + c = cases.allocate(cartesian_case, testee, "c")() + out = cases.allocate(cartesian_case, testee, cases.RETURN)() + cases.verify( + cartesian_case, + testee, + c, + out=out, + ref=tuple(t.asnumpy() + c.rho.asnumpy() for t in c.tracers), + ) + + @gtx.field_operator def constructed_inside_named_tuple( vel: tuple[gtx.Field[[IDim, JDim], gtx.float32], gtx.Field[[IDim, JDim], gtx.float32]], @@ -307,7 +346,7 @@ class StateNamedTuple(NamedTuple): value: gtx.float32 -@gtx.scan_operator(axis=cases.KDim, forward=True, init=StateDataclass(value=0.0)) +@gtx.scan_operator(axis=cases.KDim, forward=True, init=StateDataclass(value=gtx.float32(0.0))) def scan_dataclass( state: StateDataclass, inp: gtx.float32, @@ -322,7 +361,7 @@ def scan_dataclass_wrapper(inp: gtx.Field[[KDim], gtx.float32]) -> gtx.Field[[KD return scan_dataclass(inp).value -@gtx.scan_operator(axis=cases.KDim, forward=True, init=StateNamedTuple(value=0.0)) +@gtx.scan_operator(axis=cases.KDim, forward=True, init=StateNamedTuple(value=gtx.float32(0.0))) def scan_named_tuple( state: StateNamedTuple, inp: gtx.float32, diff --git a/tests/next_tests/unit_tests/type_system_tests/test_type_translation.py b/tests/next_tests/unit_tests/type_system_tests/test_type_translation.py index 6d73512150..181c1623ed 100644 --- a/tests/next_tests/unit_tests/type_system_tests/test_type_translation.py +++ b/tests/next_tests/unit_tests/type_system_tests/test_type_translation.py @@ -6,6 +6,7 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +import dataclasses import enum import types import typing @@ -313,6 +314,23 @@ def test_as_from_dtype(dtype): _FIELD = constructors.empty({cnc.TDim: 1}, gtx.float32) +_FIELD_TYPE = ts.FieldType(dims=[cnc.TDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT32)) + + +class VarLenMemberNamedTuple(typing.NamedTuple): + tracers: tuple[gtx.Field[gtx.Dims[cnc.TDim], gtx.float32], ...] + rho: gtx.Field[gtx.Dims[cnc.TDim], gtx.float32] + + +@dataclasses.dataclass +class VarLenMemberDataclass: + tracers: tuple[gtx.Field[gtx.Dims[cnc.TDim], gtx.float32], ...] + rho: gtx.Field[gtx.Dims[cnc.TDim], gtx.float32] + + +@dataclasses.dataclass +class ScalarMemberDataclass: + value: gtx.float32 @pytest.mark.parametrize( @@ -428,6 +446,42 @@ def test_from_value(value, expected): assert type_translation.from_value(value) == expected +@pytest.mark.parametrize("container", [VarLenMemberNamedTuple, VarLenMemberDataclass]) +def test_from_value_named_collection_var_len_tuple_member(container): + # the declared type keeps the variable-length tuple ... + assert type_translation.from_type_hint(container) == ts.NamedCollectionType( + types=[ts.VarLenTupleType(element_type=_FIELD_TYPE), _FIELD_TYPE], + keys=["tracers", "rho"], + original_python_type=_make_type_string_for_container(container), + ) + # ... while the member types of a value are taken from the runtime members, i.e. + # the tuple has a concrete length + value = container(tracers=(_FIELD, _FIELD), rho=_FIELD) + assert type_translation.from_value(value) == ts.NamedCollectionType( + types=[ts.TupleType(types=[_FIELD_TYPE, _FIELD_TYPE]), _FIELD_TYPE], + keys=["tracers", "rho"], + original_python_type=_make_type_string_for_container(container), + ) + + +def test_from_value_named_collection_scalar_member(): + assert type_translation.from_value( + ScalarMemberDataclass(value=np.float32(0.0)) + ) == ts.NamedCollectionType( + types=[ts.ScalarType(kind=ts.ScalarKind.FLOAT32)], + keys=["value"], + original_python_type=_make_type_string_for_container(ScalarMemberDataclass), + ) + + +@pytest.mark.skipif(not __debug__, reason="the member compatibility is only checked in debug mode") +def test_from_value_named_collection_incompatible_member(): + # a plain Python 'float' is typed 'float64' by convention, which is incompatible + # with the member annotation 'float32'; note that this is only checked in debug mode + with pytest.raises(ValueError, match="Member 'value' of 'ScalarMemberDataclass'"): + type_translation.from_value(ScalarMemberDataclass(value=0.0)) + + def test_from_value_namespace_proxy(): import next_tests.artifacts.dummy_package as dummy_package From 7bda923037ad1c943fef8a6d2e8b31278679e28d Mon Sep 17 00:00:00 2001 From: Till Ehrengruber Date: Fri, 16 Oct 2026 14:54:07 +0200 Subject: [PATCH 21/24] fix[next]: support named collections inside variable-length tuples The argument value extraction was blind to variable-length tuples: 'type_info.needs_value_extraction' did not descend into them and the extractor expression builders could only emit fixed index accesses, so e.g. a dataclass instance inside a 'tuple[SomeDataclass, ...]' argument reached the lowered program unextracted and failed at run time ('object is not subscriptable'). Extract variable-length tuples with a comprehension in both the type-spec and the type-hint based expression builder (the latter also fixes the hint form 'tuple[X, ...]', which was mis-handled as a fixed two-element tuple). Nested comprehensions reuse the variable name, relying on comprehension scoping. --- src/gt4py/next/named_collections.py | 15 ++++++ src/gt4py/next/type_system/type_info.py | 15 ++++-- .../ffront_tests/test_named_collections.py | 54 +++++++++++++++++++ .../next_tests/unit_tests/test_containers.py | 44 +++++++++++++++ .../test_type_translation.py | 15 ++++++ 5 files changed, 139 insertions(+), 4 deletions(-) diff --git a/src/gt4py/next/named_collections.py b/src/gt4py/next/named_collections.py index e2c7c09a52..31dbab9b83 100644 --- a/src/gt4py/next/named_collections.py +++ b/src/gt4py/next/named_collections.py @@ -165,6 +165,16 @@ def make_extractor_expr(named_collection_type_hint: xtyping.TypeAnnotation, valu elif issubclass(actual_type, tuple) and ( tuple_arg_hints := typing.get_args(named_collection_type_hint) ): + if len(tuple_arg_hints) == 2 and tuple_arg_hints[1] is Ellipsis: + # Variable-length tuple: the length is only known at run time, extract + # with a comprehension, or pass the value through if the elements do not + # need extraction. Note that reusing the variable name in nested + # comprehensions is fine: a comprehension's iterable is evaluated in the + # enclosing scope and the element expression only uses its own variable. + element_expr = make_extractor_expr(tuple_arg_hints[0], "__el") + if element_expr == "__el": + return value_expr + return f"tuple({element_expr} for __el in {value_expr})" # This is a `tuple` with type arguments expr_parts = [ make_extractor_expr(value, f"{value_expr}[{i}]") @@ -202,6 +212,11 @@ def make_extractor_expr_from_type_spec(type_: ts.TypeSpec, value_expr: str) -> s for i, t in enumerate(type_.types) ] return f"({', '.join(args)})" + case ts.VarLenTupleType() if type_info.needs_value_extraction(type_): + # the length is only known at run time, extract with a comprehension (see + # the note on nested comprehensions in `make_extractor_expr`) + element_expr = make_extractor_expr_from_type_spec(type_.element_type, "__el") + return f"tuple({element_expr} for __el in {value_expr})" return value_expr diff --git a/src/gt4py/next/type_system/type_info.py b/src/gt4py/next/type_system/type_info.py index 6a2080690c..4d913c2308 100644 --- a/src/gt4py/next/type_system/type_info.py +++ b/src/gt4py/next/type_system/type_info.py @@ -1050,8 +1050,15 @@ def accepts_args( def needs_value_extraction( type_spec: ts.TypeSpec, -) -> typing_extensions.TypeIs[ts.NamedCollectionType | ts.TupleType]: - return isinstance(type_spec, ts.NamedCollectionType) or ( - isinstance(type_spec, ts.TupleType) - and any(needs_value_extraction(t) for t in type_spec.types) +) -> typing_extensions.TypeIs[ts.NamedCollectionType | ts.TupleType | ts.VarLenTupleType]: + return ( + isinstance(type_spec, ts.NamedCollectionType) + or ( + isinstance(type_spec, ts.TupleType) + and any(needs_value_extraction(t) for t in type_spec.types) + ) + or ( + isinstance(type_spec, ts.VarLenTupleType) + and needs_value_extraction(type_spec.element_type) + ) ) diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_named_collections.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_named_collections.py index 302f03c63b..3ef33405fd 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_named_collections.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_named_collections.py @@ -149,6 +149,60 @@ def test_named_collection_with_var_len_tuple_member(cartesian_case, testee): ) +@gtx.field_operator +def sum_named_tuple_elements( + pairs: tuple[NamedTupleNamedCollection, ...], +) -> tuple[gtx.Field[[IDim, JDim], gtx.float32], ...]: + return tuple(p.u + p.v for p in pairs) + + +@gtx.field_operator +def sum_dataclass_elements( + pairs: tuple[DataclassNamedCollection, ...], +) -> tuple[gtx.Field[[IDim, JDim], gtx.float32], ...]: + return tuple(p.u + p.v for p in pairs) + + +@pytest.mark.parametrize("testee", [sum_named_tuple_elements, sum_dataclass_elements]) +@pytest.mark.uses_tuple_args +@pytest.mark.uses_tuple_returns +def test_var_len_tuple_of_named_collections(cartesian_case, testee): + pairs = cases.allocate(cartesian_case, testee, "pairs")() + out = cases.allocate(cartesian_case, testee, cases.RETURN)() + cases.verify( + cartesian_case, + testee, + pairs, + out=out, + ref=tuple(p.u.asnumpy() + p.v.asnumpy() for p in pairs), + ) + + +@dataclasses.dataclass +class PairsContainer: + pairs: tuple[DataclassNamedCollection, ...] + offset: gtx.Field[[IDim, JDim], gtx.float32] + + +@gtx.field_operator +def sum_nested_pairs(c: PairsContainer) -> tuple[gtx.Field[[IDim, JDim], gtx.float32], ...]: + return tuple(p.u + p.v + c.offset for p in c.pairs) + + +@pytest.mark.uses_tuple_args +@pytest.mark.uses_tuple_returns +def test_named_collection_with_var_len_tuple_of_named_collections_member(cartesian_case): + c = cases.allocate(cartesian_case, sum_nested_pairs, "c")() + out = cases.allocate(cartesian_case, sum_nested_pairs, cases.RETURN)() + cases.verify( + cartesian_case, + sum_nested_pairs, + c, + out=out, + ref=tuple(p.u.asnumpy() + p.v.asnumpy() + c.offset.asnumpy() for p in c.pairs), + ) + + @gtx.field_operator def constructed_inside_named_tuple( vel: tuple[gtx.Field[[IDim, JDim], gtx.float32], gtx.Field[[IDim, JDim], gtx.float32]], diff --git a/tests/next_tests/unit_tests/test_containers.py b/tests/next_tests/unit_tests/test_containers.py index 8704174b09..4cd3e53e24 100644 --- a/tests/next_tests/unit_tests/test_containers.py +++ b/tests/next_tests/unit_tests/test_containers.py @@ -72,6 +72,50 @@ def test_make_container_extractor( ) +@dataclasses.dataclass +class PointDataclass: + x: int + y: int + + +@dataclasses.dataclass +class VarLenPointsDataclass: + points: tuple[PointDataclass, ...] + offset: int + + +def test_make_container_extractor_var_len_tuple(): + # the length of the variable-length tuple member is only known from the value, the + # same extractor must handle all lengths + extractor = named_collections.make_named_collection_extractor(VarLenPointsDataclass) + value = VarLenPointsDataclass( + points=(PointDataclass(1, 2), PointDataclass(3, 4), PointDataclass(5, 6)), offset=7 + ) + assert extractor(value) == (((1, 2), (3, 4), (5, 6)), 7) + assert extractor(VarLenPointsDataclass(points=(PointDataclass(1, 2),), offset=0)) == ( + ((1, 2),), + 0, + ) + + +@dataclasses.dataclass +class NestedVarLenDataclass: + groups: tuple[VarLenPointsDataclass, ...] + + +def test_make_container_extractor_nested_var_len_tuple(): + # nested comprehensions in the generated extractor reuse the variable name, + # relying on comprehension scoping + extractor = named_collections.make_named_collection_extractor(NestedVarLenDataclass) + value = NestedVarLenDataclass( + groups=( + VarLenPointsDataclass(points=(PointDataclass(1, 2), PointDataclass(3, 4)), offset=5), + VarLenPointsDataclass(points=(PointDataclass(6, 7),), offset=8), + ) + ) + assert extractor(value) == (((((1, 2), (3, 4)), 5), (((6, 7),), 8)),) + + @pytest.mark.parametrize( ["nested_tuple", "expected_container"], [(value, cnc.from_nested_tuple(cls, value)) for cls, value in cnc.PYCONTAINERS_SAMPLES.items()], diff --git a/tests/next_tests/unit_tests/type_system_tests/test_type_translation.py b/tests/next_tests/unit_tests/type_system_tests/test_type_translation.py index 181c1623ed..7244545587 100644 --- a/tests/next_tests/unit_tests/type_system_tests/test_type_translation.py +++ b/tests/next_tests/unit_tests/type_system_tests/test_type_translation.py @@ -464,6 +464,21 @@ def test_from_value_named_collection_var_len_tuple_member(container): ) +def test_from_value_var_len_tuple_of_named_collections(): + # the declared type is a variable-length tuple of named collections ... + assert type_translation.from_type_hint( + tuple[cnc.NamedTupleNamedCollection, ...] + ) == ts.VarLenTupleType(element_type=NamedTupleNamedCollectionTypeSpec) + # ... while the type of a value has a concrete length + value = ( + cnc.NamedTupleNamedCollection(x=_FIELD, y=_FIELD), + cnc.NamedTupleNamedCollection(x=_FIELD, y=_FIELD), + ) + assert type_translation.from_value(value) == ts.TupleType( + types=[NamedTupleNamedCollectionTypeSpec, NamedTupleNamedCollectionTypeSpec] + ) + + def test_from_value_named_collection_scalar_member(): assert type_translation.from_value( ScalarMemberDataclass(value=np.float32(0.0)) From e2d872e23eb3ca0de97b2f8e8b37bb28bf429cc5 Mon Sep 17 00:00:00 2001 From: Till Ehrengruber Date: Fri, 16 Oct 2026 16:51:48 +0200 Subject: [PATCH 22/24] fix[next]: handle variable-length tuple arguments of scan operators Direct calls to scan operators with a variable-length tuple parameter crashed when building the compiled-programs pool and in the scan signature check: - Allow 'DeferredType' as the element type of a 'VarLenTupleType' (to be removed again together with 'DeferredType' in 'TupleType.types') so the structure-preserving deferred program-context type of a scan operator can represent variable-length tuple parameters. - Make the scan parameter promotion variable-length tuple aware: a concrete tuple argument uses its first element as representative (requiring all elements to be equal), a still generic argument (e.g. a variable-length tuple program parameter forwarded to the scan) uses the element type. Note that such arguments can not be consumed in the scan body yet (subscripts and reductions on variable-length tuples do not exist) and the carry must have a fixed length as the 'init' value is necessarily concrete. --- .../ffront/foast_passes/type_deduction.py | 1 + src/gt4py/next/ffront/type_info.py | 26 +++++--- src/gt4py/next/type_system/type_info.py | 2 +- .../next/type_system/type_specifications.py | 3 +- .../ffront_tests/test_func_to_past.py | 18 ++++++ .../type_system_tests/test_type_info.py | 62 +++++++++++++++++++ 6 files changed, 102 insertions(+), 10 deletions(-) diff --git a/src/gt4py/next/ffront/foast_passes/type_deduction.py b/src/gt4py/next/ffront/foast_passes/type_deduction.py index 79bdcac56e..f21c85f2c4 100644 --- a/src/gt4py/next/ffront/foast_passes/type_deduction.py +++ b/src/gt4py/next/ffront/foast_passes/type_deduction.py @@ -865,6 +865,7 @@ def visit_TupleComprehension( return result elif isinstance(iterable.type, ts.VarLenTupleType): element_type = iterable.type.element_type + assert isinstance(element_type, ts.DataType) # 'DeferredType' can not occur here new_mapper = self._deduce_tuple_comprehension_mapper( node, target, element_type, **kwargs ) diff --git a/src/gt4py/next/ffront/type_info.py b/src/gt4py/next/ffront/type_info.py index ffdf7b63bf..e2e455387c 100644 --- a/src/gt4py/next/ffront/type_info.py +++ b/src/gt4py/next/ffront/type_info.py @@ -164,7 +164,7 @@ def _tree_map_type_constructor_drop_python_type( def _scan_param_promotion( param: ts.TypeSpec, arg: ts.TypeSpec -) -> ts.FieldType | ts.TupleType | ts.NamedCollectionType: +) -> ts.FieldType | ts.TupleType | ts.NamedCollectionType | ts.VarLenTupleType: """ Promote parameter of a scan pass to match dimensions of respective scan operator argument. @@ -183,14 +183,24 @@ def _scan_param_promotion( FieldType(dims=[Dimension(value='I', kind=)], dtype=ScalarType(kind=, shape=None)) """ - def _as_field(dtype: ts.TypeSpec, path: tuple[int, ...]) -> ts.FieldType: + def _index_arg(type_: ts.TypeSpec, idx: int | type_info.VarLenTupleElementKey) -> ts.TypeSpec: + if isinstance(idx, type_info.VarLenTupleElementKey): + if isinstance(type_, ts.TupleType): + # case where `param` is variable length, but `arg` is actual tuple + if not all(el == type_.types[0] for el in type_.types): # type: ignore[attr-defined] + raise IndexError() + return type_.types[0] # type: ignore[attr-defined] + elif isinstance(type_, ts.VarLenTupleType): + return type_.element_type + raise AttributeError + return type_.types[idx] # type: ignore[attr-defined] + + def _as_field( + dtype: ts.TypeSpec, path: tuple[int | type_info.VarLenTupleElementKey, ...] + ) -> ts.FieldType: assert isinstance(dtype, ts.ScalarType) try: - el_type = functools.reduce( - lambda type_, idx: type_.types[idx], # type: ignore[attr-defined] - path, - arg, - ) + el_type = functools.reduce(_index_arg, path, arg) return ts.FieldType(dims=type_info.extract_dims(el_type), dtype=dtype) except (IndexError, AttributeError): # The structure of the scan passes argument and the requested @@ -208,7 +218,7 @@ def _as_field(dtype: ts.TypeSpec, path: tuple[int, ...]) -> ts.FieldType: result_collection_constructor=_tree_map_type_constructor_drop_python_type, with_path_arg=True, )(param) - assert isinstance(res, (ts.FieldType, ts.TupleType, ts.NamedCollectionType)) + assert isinstance(res, (ts.FieldType, ts.TupleType, ts.NamedCollectionType, ts.VarLenTupleType)) return res diff --git a/src/gt4py/next/type_system/type_info.py b/src/gt4py/next/type_system/type_info.py index 4d913c2308..2d616f6c5c 100644 --- a/src/gt4py/next/type_system/type_info.py +++ b/src/gt4py/next/type_system/type_info.py @@ -180,7 +180,7 @@ def tree_map_type_constructor( ) -> ts.CollectionTypeSpecT | ts.VarLenTupleType: if isinstance(value, ts.VarLenTupleType): (element_type,) = elems - assert isinstance(element_type, ts.DataType) + assert isinstance(element_type, (ts.DataType, ts.DeferredType)) return ts.VarLenTupleType(element_type=element_type) return ( ts.NamedCollectionType( diff --git a/src/gt4py/next/type_system/type_specifications.py b/src/gt4py/next/type_system/type_specifications.py index 91edc45457..3261ab0e36 100644 --- a/src/gt4py/next/type_system/type_specifications.py +++ b/src/gt4py/next/type_system/type_specifications.py @@ -150,7 +150,8 @@ class VarLenTupleType(DataType): `TupleType` (or stays variable-length in length-agnostic code paths). """ - element_type: DataType + # TODO(tehrengruber): `DeferredType` should be removed again, see `TupleType.types`. + element_type: DataType | DeferredType def __str__(self) -> str: return f"tuple[{self.element_type}, ...]" diff --git a/tests/next_tests/unit_tests/ffront_tests/test_func_to_past.py b/tests/next_tests/unit_tests/ffront_tests/test_func_to_past.py index 533113bfc2..2c895cc9e2 100644 --- a/tests/next_tests/unit_tests/ffront_tests/test_func_to_past.py +++ b/tests/next_tests/unit_tests/ffront_tests/test_func_to_past.py @@ -228,6 +228,24 @@ def var_len_tuple_domain_program( ) +def test_var_len_tuple_scan_arg(): + # A variable-length tuple program parameter forwarded to a scan operator: the scan + # signature check runs against the still generic (declared) argument type. + KDim = gtx.Dimension("KDim", kind=gtx.DimensionKind.VERTICAL) + + @gtx.scan_operator(axis=KDim, forward=True, init=float64(0.0)) + def scan_tracers(carry: float64, tracers: tuple[float64, ...]) -> float64: + return carry + + def scan_program( + tracers: tuple[gtx.Field[[KDim], float64], ...], + out: gtx.Field[[KDim], float64], + ): + scan_tracers(tracers, out=out) + + ProgramParser.apply_to_function(scan_program) + + def test_domain_exception_1(identity_def): domain_format_1 = gtx.field_operator(identity_def) diff --git a/tests/next_tests/unit_tests/type_system_tests/test_type_info.py b/tests/next_tests/unit_tests/type_system_tests/test_type_info.py index 1d52ba2e38..be7f3aa659 100644 --- a/tests/next_tests/unit_tests/type_system_tests/test_type_info.py +++ b/tests/next_tests/unit_tests/type_system_tests/test_type_info.py @@ -117,6 +117,18 @@ def callable_type_info_cases(): returns=float_type, ), ) + var_len_scanop_type = ts_ffront.ScanOperatorType( + axis=KDim, + definition=ts.FunctionType( + pos_only_args=[], + pos_or_kw_args={ + "carry": float_type, + "a": ts.VarLenTupleType(element_type=int_type), + }, + kw_only_args={}, + returns=float_type, + ), + ) return [ # func_type, pos_only_args, kwargs, expected incompatibilities, return type @@ -348,6 +360,56 @@ def callable_type_info_cases(): ], ts.FieldType(dims=[IDim, JDim, KDim], dtype=float_type), ), + # scan operator with a variable-length tuple parameter + ( + # concrete, homogeneous argument + var_len_scanop_type, + [ + ts.TupleType( + types=[ + ts.FieldType(dims=[IDim, JDim], dtype=int_type), + ts.FieldType(dims=[IDim, JDim], dtype=int_type), + ] + ) + ], + {}, + [], + ts.FieldType(dims=[IDim, JDim, KDim], dtype=float_type), + ), + ( + # still generic argument, e.g. a variable-length tuple parameter of a program + var_len_scanop_type, + [ts.VarLenTupleType(element_type=ts.FieldType(dims=[IDim, JDim], dtype=int_type))], + {}, + [], + ts.FieldType(dims=[IDim, JDim, KDim], dtype=float_type), + ), + ( + # heterogeneous argument + var_len_scanop_type, + [ + ts.TupleType( + types=[ + ts.FieldType(dims=[IDim], dtype=int_type), + ts.FieldType(dims=[JDim], dtype=int_type), + ] + ) + ], + {}, + [ + r"Expected argument 'a' to be of type 'tuple\[Field\[\[\.\.\.\], int64\], " + r"\.\.\.\]', got 'tuple\[Field\[\[I\], int64\], Field\[\[J\], int64\]\]'." + ], + ts.FieldType(dims=[IDim, JDim, KDim], dtype=float_type), + ), + ( + # the empty tuple is a valid concretization of a variable-length tuple + var_len_scanop_type, + [ts.TupleType(types=[])], + {}, + [], + ts.FieldType(dims=[KDim], dtype=float_type), + ), ( ts.FunctionType( pos_only_args=[ From 78bdccb47edccff45ed5f0d9ca928ea7f750e4de Mon Sep 17 00:00:00 2001 From: Till Ehrengruber Date: Fri, 16 Oct 2026 16:56:08 +0200 Subject: [PATCH 23/24] fix[next]: proper error for static domains with variable-length tuples The elements of a variable-length tuple parameter can not be addressed in the static-domain argument descriptors before the concrete length is known at the call; reject such programs with a proper error message instead of a bare 'NotImplementedError'. --- src/gt4py/next/ffront/decorator.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/gt4py/next/ffront/decorator.py b/src/gt4py/next/ffront/decorator.py index c94341b7da..e13a85b9b2 100644 --- a/src/gt4py/next/ffront/decorator.py +++ b/src/gt4py/next/ffront/decorator.py @@ -215,7 +215,11 @@ def _field_domain_descriptor_mapping_from_func_type(func_type: ts.FunctionType) for el_type_, path in type_info.primitive_constituents(type_, with_path_arg=True): if isinstance(el_type_, ts.FieldType): if any(isinstance(idx, type_info.VarLenTupleElementKey) for idx in path): - raise NotImplementedError() + raise NotImplementedError( + f"Static domains are not supported for parameter '{name}': the " + "elements of a variable-length tuple can not be addressed before " + "the concrete length is known at the call." + ) path_as_expr = "".join(f"[{idx}]" for idx in path) static_domain_args.append(f"{name}{path_as_expr}") return static_domain_args From a8dbef6fe8c08374a80682b0a03650e9e4e29eff Mon Sep 17 00:00:00 2001 From: Till Ehrengruber Date: Fri, 16 Oct 2026 17:05:59 +0200 Subject: [PATCH 24/24] style[next]: fix code quality checks Missing blank line and two 'type: ignore' comments that became unnecessary with the 'TupleType' narrowing in '_index_arg'. --- src/gt4py/next/ffront/type_info.py | 4 ++-- .../unit_tests/ffront_tests/test_diagnostic_messages.py | 1 + 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/src/gt4py/next/ffront/type_info.py b/src/gt4py/next/ffront/type_info.py index e2e455387c..b26cea2cc7 100644 --- a/src/gt4py/next/ffront/type_info.py +++ b/src/gt4py/next/ffront/type_info.py @@ -187,9 +187,9 @@ def _index_arg(type_: ts.TypeSpec, idx: int | type_info.VarLenTupleElementKey) - if isinstance(idx, type_info.VarLenTupleElementKey): if isinstance(type_, ts.TupleType): # case where `param` is variable length, but `arg` is actual tuple - if not all(el == type_.types[0] for el in type_.types): # type: ignore[attr-defined] + if not all(el == type_.types[0] for el in type_.types): raise IndexError() - return type_.types[0] # type: ignore[attr-defined] + return type_.types[0] elif isinstance(type_, ts.VarLenTupleType): return type_.element_type raise AttributeError diff --git a/tests/next_tests/unit_tests/ffront_tests/test_diagnostic_messages.py b/tests/next_tests/unit_tests/ffront_tests/test_diagnostic_messages.py index 2b899e4304..4ace4f8913 100644 --- a/tests/next_tests/unit_tests/ffront_tests/test_diagnostic_messages.py +++ b/tests/next_tests/unit_tests/ffront_tests/test_diagnostic_messages.py @@ -393,6 +393,7 @@ def with_filtered_genexp(a: gtx.Field[[IDim], float64]) -> gtx.Field[[IDim], flo assert any("known at compile time" in note for note in err.notes) assert "if True" in str(err) + def test_unindexed_cartesian_offset_names_the_offset_as_written(): # The tag of 'IOff' is 'Ioff'; the message has to quote what the user wrote. def unindexed(a: gtx.Field[[IDim], float64]) -> gtx.Field[[IDim], float64]: